update folder

data_setup.py

@@ -0,0 +1,137 @@
+"""
+contains functionality for creating pytorch dataloaders for image classification data
+"""
+import os 
+import torch
+from torchvision import datasets, transforms 
+from torch.utils.data import DataLoader 
+from pathlib import Path 
+import pathlib
+import requests
+import zipfile
+from typing import Tuple, Dict, List
+from torch.utils.data import Dataset
+from PIL import Image
+
+NUM_WORKERS = os.cpu_count()
+
+# create custom dataset
+def find_classes(directory: str) -> Tuple[list[str], Dict[str, int]]:
+    """
+    Finds the class folder names in a target directory 
+    """
+    # 1. get the class names by scanning the target directory 
+    classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())
+
+    # 2. raise an error is class names couldn't be found 
+    if not classes:
+        raise FileNotFoundError(f"couldn't find any classes in {directory}")
+    
+    # 3. create a dictionary of index labels 
+    class_to_idx = {class_name: i for i, class_name in enumerate(classes)}
+    return classes, class_to_idx
+
+# 1. subclass torch.utils.data.Dataset 
+class ImageFolderCustom(Dataset):
+    # 2. initialize the constructor
+    def __init__(self, targ_dir: str, heads: list[str], transform=None, is_training: bool = True):
+        # 3. create several attributes 
+        # get all the image paths
+        self.training = []
+        self.testing = []
+        for tag in heads: 
+            self.img_list = list(Path(targ_dir / tag).glob("*.jpg"))
+            self.train_length = int(len(self.img_list) * 0.8)
+            self.training.extend(self.img_list[:self.train_length])
+            self.testing.extend(self.img_list[self.train_length:])
+
+        if is_training: 
+            self.paths = self.training
+        else: 
+            self.paths = self.testing
+        # setup transforms
+        self.transform = transform
+        # create classes and class_to_idx 
+        self.classes, self.class_to_idx = find_classes(targ_dir)
+
+    # 4. create a function to load images 
+    def load_image(self, index: int) -> Image.Image: 
+        "opens an image via a path and returns it"
+        image_path = self.paths[index]
+        return Image.open(image_path)
+    
+    # 5. overwrite __len__()
+    def __len__(self) -> int: 
+        return len(self.paths)
+    
+    # 6. overwrite __getitem__() to return a particular sample
+    def __getitem__(self, index: int) -> Tuple[torch.Tensor, int]:
+        "returns one sample of data, data and the label (X, y)"
+        img = self.load_image(index)
+        class_name = self.paths[index].parent.name # expects path in format: data_folder/class_name/image.jpg
+        class_idx = self.class_to_idx[class_name]
+
+        # transform if necessary 
+        if self.transform:
+            return self.transform(img), class_idx
+        else: 
+            return img, class_idx
+
+def create_dataloaders(
+    image_dir: str,  
+    heads: list[str],
+    train_transform: transforms.Compose, 
+    test_transform: transforms.Compose,
+    batch_size: int, 
+    num_workers: int=NUM_WORKERS
+):
+    """
+    creates training and testing DataLoaders. 
+
+    Takes in a training directory and testing directory path and turns them
+    into pytorch datasets and then into pytorch dataloaders. 
+
+    Args:
+        train_dir: path to training directory. 
+        test_dir: path to testing directory 
+        transform: torchvision transforms to perform on training and testing data. 
+        batch_size: number of samples per batch in each of the dataloaders. 
+        num_workers: an integer for number of workers per dataloader.
+
+    returns: 
+        A tuple of (train_dataloader, test_dataloader, class_names).
+        where class_names is a list of the target classes. 
+
+        Example usage: 
+            train_dataloader, test_dataloader, class_names = create_dataloaders(train_dir=path/to/train_dir, 
+                                                                                test_dir=path/to/test_dir, 
+                                                                                transform=some_transform,
+                                                                                batch_size=32,
+                                                                                num_workers=4)                                                                                                                 
+    """
+
+    # use ImageFolder to create datasets 
+    train_data = ImageFolderCustom(targ_dir=image_dir, heads=heads, transform=train_transform, is_training=True)
+
+    test_data = ImageFolderCustom(targ_dir=image_dir, heads=heads, transform=test_transform, is_training=False)
+
+    # get class names 
+    class_names = train_data.classes 
+
+    # turn images into dataloaders 
+    train_dataloader = DataLoader(
+        train_data, 
+        batch_size=batch_size, 
+        shuffle=True, 
+        num_workers=num_workers, 
+        pin_memory=True
+    )
+    test_dataloader = DataLoader(
+        test_data, 
+        batch_size=batch_size, 
+        shuffle=False, 
+        num_workers=num_workers, 
+        pin_memory=True
+    )
+
+    return train_dataloader, test_dataloader, class_names

engine.py

@@ -0,0 +1,194 @@
+"""
+contains functions for training and testing a pytorch model
+"""
+import torch 
+
+from tqdm.auto import tqdm 
+from typing import Dict, List, Tuple 
+# from torch.utils.tensorboard.writer import SummaryWriter
+
+def train_step(model: torch.nn.Module, 
+              dataloader: torch.utils.data.DataLoader, 
+              loss_fn: torch.nn.Module, 
+              optimizer: torch.optim.Optimizer, 
+              device: torch.device) -> Tuple[float, float]:
+    """Trains a pytorch model for a single epoch 
+
+    turns a target model to training mode then runs through all of the required training steps
+    (forward pass, loss calculation, optimizer step).
+
+    Args: 
+        model: pytorch model
+        dataloader: dataloader insatnce for the model to be trained on 
+        loss_fn: pytorch loss function to calculate loss
+        optimizer: pytorch optimizer to help minimize the loss function
+        device: target device
+
+    returns:
+        a tuple of training loss and training accuracy metrics
+        in the form (train_loss, train_accuracy)
+    """
+    # put the model into training mode
+    model.train()
+    
+    # setup train loss and train accuracy 
+    train_loss, train_accuracy = 0, 0 
+
+    # loop through data laoder batches
+    for batch, (X, y) in enumerate(dataloader):
+        # send data to target device 
+        X, y = X.to(device), y.to(device)
+
+        # forward pass 
+        logits = model(X)
+
+        # calculate loss and accumulate loss 
+        loss = loss_fn(logits, y)
+        train_loss += loss
+
+        # optimizer zero grad 
+        optimizer.zero_grad()
+
+        # loss backward 
+        loss.backward()
+
+        # optimizer step 
+        optimizer.step()
+
+        # calculate and accumulate accuracy metric across all batches
+        preds = torch.softmax(logits, dim=-1).argmax(dim=-1)
+        train_accuracy += (preds == y).sum().item()/len(preds)
+
+    # adjust metrics to get average loss and accuracy per batch 
+    train_loss /= len(dataloader)
+    train_accuracy /= len(dataloader)
+    return train_loss, train_accuracy
+
+def test_step(model: torch.nn.Module, 
+             dataloader: torch.utils.data.DataLoader, 
+             loss_fn: torch.nn.Module, 
+             device: torch.device) -> Tuple[float, float]:
+    """Tests a pytorch model for a single epoch
+
+    Turns a target model to eval mode and then performs a forward pass on a testing
+    dataset. 
+
+    Args: 
+        model: pytorch model
+        dataloader: dataloader insatnce for the model to be tested on 
+        loss_fn: loss function to calculate loss (errors)
+        device: target device to compute on 
+
+    returns:
+        A tuple of testing loss and testing accuracy metrics.
+        In the form (test_loss, test_accuracy)
+    """
+    # put the model in eval mode
+    model.eval()
+
+    # setup test loss and test accuracy 
+    test_loss, test_accuracy = 0, 0 
+
+    # turn on inference mode 
+    with torch.inference_mode():
+        # loop through all batches 
+        for X, y in dataloader: 
+            # send data to target device
+            X, y  = X.to(device), y.to(device)
+
+            # forward pass
+            logits = model(X)
+
+            # calculate and accumulate loss
+            loss = loss_fn(logits, y)
+            test_loss += loss.item()
+
+            # calculate and accumulate accuracy 
+            test_preds = torch.softmax(logits, dim=-1).argmax(dim=-1)
+            test_accuracy += ((test_preds == y).sum().item()/len(test_preds))
+    # adjust metrics to get average loss and accuracy per batch 
+    test_loss /= len(dataloader)
+    test_accuracy /= len(dataloader)
+    return test_loss, test_accuracy
+
+def train(model: torch.nn.Module, 
+         train_dataloader: torch.utils.data.DataLoader, 
+         test_dataloader: torch.utils.data.DataLoader, 
+         optimizer: torch.optim.Optimizer, 
+         loss_fn: torch.nn.Module, 
+         epochs: int, 
+         device: torch.device, 
+         writer: torch.utils.tensorboard.writer.SummaryWriter) -> Dict[str, List]:
+    """Trains and tests pytorch model
+
+    passes a target model through train_step() and test_step() 
+    functions for a number of epochs, training and testing the model in the same epoch loop.
+
+    calculates, prints and stores evaluation metric throughout. 
+
+    Args: 
+        model: pytorch model
+        train_dataloader: DataLoader instance for the model to be trained on
+        test_dataloader: DataLoader instance for the model to be tested on
+        optimizer: pytorch optimizer
+        loss_fn: pytorch loss function
+        epochs: integer indicating how many epochs to train for
+        device: target device to compute on 
+
+    returns: 
+        A dictionaru of training and testing loss as well as training and testing accuracy 
+        metrics. Each metric has a value in a list for each epoch. 
+
+        In the form: {train_loss: [...],
+                      train_acc: [...],
+                      test_loss: [...],
+                      test_acc: [...]}
+    """
+    # create an empty dictionary 
+    results = {
+        "train_loss": [],
+        "train_acc": [],
+        "test_loss": [],
+        "test_acc": []
+    }
+
+    # loop through training and testing steps for a number of epochs 
+    for epoch in tqdm(range(epochs)):
+        train_loss, train_acc = train_step(model=model,
+                                          dataloader=train_dataloader, 
+                                          loss_fn=loss_fn, 
+                                          optimizer=optimizer, 
+                                          device=device)
+        test_loss, test_acc = test_step(model=model, 
+                                       dataloader=test_dataloader, 
+                                       loss_fn=loss_fn,
+                                       device=device)
+
+        if epoch % 1 == 0:
+            print(
+                f"Epoch: {epoch+1} | " 
+                f"train_loss: {train_loss:.4f} | "
+                f"train_acc: {train_acc:.4f} | "
+                f"test_loss: {test_loss:.4f} | "
+                f"test_acc: {test_acc:.4f}"
+            )
+
+        # update results dictionary 
+        results["train_loss"].append(train_loss.item())
+        results["train_acc"].append(train_acc)
+        results["test_loss"].append(test_loss)
+        results["test_acc"].append(test_acc)
+
+        if writer: 
+            # NEW: EXPERIMENT TRACKING 
+            # add loss to SummaryWriter
+            writer.add_scalars(main_tag="Loss", tag_scalar_dict={"train loss": train_loss, "test loss": test_loss}, global_step=epoch)
+            # add accuracy to SummaryWriter 
+            writer.add_scalars(main_tag="Accuracy", tag_scalar_dict={"train acc": train_acc, "test acc": test_acc}, global_step=epoch)
+            # track the pytorch model architecture 
+            writer.add_graph(model=model, input_to_model=torch.randn(size=(32, 3, 224, 224)).to(device))
+            writer.close()
+    # END SummaryWriter tracking process
+
+    # return the filled results dictionaru 
+    return results

experiments.py

@@ -0,0 +1,41 @@
+import torch
+import utils
+import model_builder as mb
+import engine
+
+# 3. loop through each dataloader 
+def run_experiment(train_dataloaders: dict, test_dataloader: torch.utils.data.DataLoader,  num_epochs: int, models: list[str], class_names: list[str], device: torch.device = None):
+    # 1. set seed
+    utils.set_seeds(seed=42)
+
+    # 2. keep track of experiment numbers 
+    experiment_number = 0 
+    for dataloader_name, train_dataloader in train_dataloaders.items():
+    # 4. loop through each number of epochs 
+        for epochs in num_epochs:
+        # 5. loop through each model name and create a new model based on the name 
+            for model_name in models:
+            # 6. create information prints out 
+                experiment_number += 1 
+                print(f"[INFO] experiment number: {experiment_number}")
+                print(f"[INFO] model: {model_name}")
+                print(f"[INFO] dataloader: {dataloader_name}")
+                print(f"[INFO] number of epochs: {epochs}")
+
+                # 7. select the model 
+                if model_name == "effnetb0":
+                    model = mb.create_model_baseline_effnetb0(out_feats=len(class_names), device=device)
+                else: 
+                    model = mb.create_model_baseline_effnetb2(out_feats=len(class_names), device=device)
+            
+                # 8. create a new loss function for every model 
+                loss_fn = torch.nn.CrossEntropyLoss()
+                optimizer = torch.optim.Adam(params=model.parameters(), lr=0.001)
+
+                # 9. train target model with target dataloaders and track experiment
+                engine.train(model=model, train_dataloader=train_dataloader, test_dataloader=test_dataloader, optimizer=optimizer, loss_fn=loss_fn, epochs=epochs, device=device, writer=utils.create_writer(experiment_name=dataloader_name, model_name=model_name, extra=f"{epochs}_epochs"))
+
+                # 10. save the model to file
+                save_filepath = f"{model_name}_{dataloader_name}_{epochs}_epochs.pt"
+                utils.save_model(model=model, target_dir="models", model_name=save_filepath)
+                print("-"*50+"\n")

model_builder.py

@@ -0,0 +1,50 @@
+"""
+contains pytorch model code to instantiate a TinyVGG model.
+"""
+import torch 
+from torch import nn 
+import torchvision
+
+def create_model_baseline_effnetb0(out_feats: int, device: torch.device = None) -> torch.nn.Module:
+    weights = torchvision.models.EfficientNet_B0_Weights.DEFAULT
+    model = torchvision.models.efficientnet_b0(weights=weights).to(device)
+
+    for param in model.features.parameters():
+        param.requires_grad = False
+
+    torch.manual_seed(42)
+    torch.cuda.manual_seed(42)
+
+    # change the output layer 
+    model.classifier = torch.nn.Sequential(
+        torch.nn.Dropout(p=0.2, inplace=True), 
+        torch.nn.Linear(in_features=1280, 
+                        out_features=out_feats, 
+                        bias=True)).to(device)
+    
+    model.name = "effnetb0"
+    print(f"[INFO] created a model {model.name}")
+    
+    return model
+
+def create_model_baseline_effnetb2(out_feats: int, device: torch.device = None) -> torch.nn.Module: 
+    weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT 
+    model = torchvision.models.efficientnet_b2(weights=weights).to(device)
+
+    for param in model.features.parameters():
+        param.requires_grad = False 
+
+    torch.manual_seed(42)
+    torch.cuda.manual_seed(42)
+
+    model.classifier = nn.Sequential(
+        nn.Dropout(p=0.3, inplace=True), 
+        nn.Linear(in_features=1408,
+                  out_features=out_feats, 
+                  bias=True)
+    ).to(device)
+
+    model.name = "effnetb2"
+    print(f"[INFO] created a model {model.name}")
+    
+    return model

predict.py

@@ -0,0 +1,49 @@
+import argparse
+import torch 
+import matplotlib.pyplot as plt 
+import requests
+from PIL import Image
+from torchvision import transforms
+import data_setup, model_builder
+from pathlib import Path
+import os
+
+parser = argparse.ArgumentParser()
+parser.add_argument("-i", "--image", help="string of url to the image", type=str)
+args = parser.parse_args()
+
+URL = args.image # required
+
+image_transform = transforms.Compose([
+            transforms.Resize(size=(224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                std=[0.229, 0.224, 0.225])])
+
+IMAGE_PATH = Path("data") / "spoiled-fresh" / "FRUIT-16K"
+
+classes = sorted(entry.name for entry in os.scandir(IMAGE_PATH) if entry.is_dir())
+
+# load saved model 
+loaded_model = model_builder.create_model_baseline_effnetb2(out_feats=len(classes), device="cpu")
+loaded_model.load_state_dict(torch.load("models/effnetb2_fruitsvegs0_5_epochs.pt", weights_only=True))
+
+def pred_and_plot(model: torch.nn.Module, 
+                    image_path: str,
+                    transform: transforms.Compose,
+                    class_names: list[str] = None):
+        # load image
+        img = Image.open(requests.get(image_path, stream=True).raw).convert("RGB")
+        # setup transformed image 
+        transformed_img = transform(img)
+        # forward pass 
+        logits = model(transformed_img.unsqueeze(dim=0))
+        pred = torch.softmax(logits, dim=-1).argmax(dim=-1)
+        # plot the image along with the label 
+        # plt.imshow(transformed_img.permute(1, 2, 0))
+        title = f"{class_names[pred]} | {torch.softmax(logits, dim=-1).max():.3f}"
+        plt.title(title)
+        print(title)
+
+pred_and_plot(model=loaded_model, image_path=URL, 
+                transform=image_transform, class_names=classes)

train.py

@@ -0,0 +1,66 @@
+import os 
+import torch 
+import data_setup, engine, model_builder, utils
+from torchvision import transforms, models
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument("-e", "--num_epochs", help="an integer to perform number of epochs", type=int)
+parser.add_argument("-b", "--batch_size", help="an integer of number of element per batch", type=int)
+# parser.add_argument("-hu", "--hidden_units", help="an integer of number of hidden units per layer", type=int)
+parser.add_argument("-lr", "--learning_rate", help="a float for the learning rate", type=float)
+
+args = parser.parse_args()
+
+# setup hyperparameters 
+NUM_EPOCHS = args.num_epochs if args.num_epochs else 10
+BATCH_SIZE = args.batch_size # required
+# HIDDEN_UNITS = args.hidden_units if args.hidden_units else 10
+LEARNING_RATE = args.learning_rate if args.learning_rate else 0.001
+
+# setup directories 
+train_dir = "data/pizza_sushi_steak/train"
+test_dir = "data/pizza_sushi_steak/test"
+
+def main():
+    # setup device agnostic code 
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # create transforms
+    data_transform = transforms.Compose([
+            transforms.Resize(size=(224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                std=[0.229, 0.224, 0.225]),
+        ])
+
+    # create DataLoaders with help from data_setup.py
+    train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(
+        train_dir=train_dir, 
+        test_dir=test_dir, 
+        transform=data_transform, 
+        batch_size=BATCH_SIZE,
+        num_workers=0
+    )
+
+    # create model with help from model_builder.py
+    model = model_builder.create_model_baseline_effnetb0(out_feats=len(class_names), device=device)
+
+    # set loss and optimizer
+    loss_fn = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(params=model.parameters(), lr=LEARNING_RATE)
+
+    # start training with help from engine.py 
+    engine.train(model=model, 
+                    train_dataloader=train_dataloader, 
+                    test_dataloader=test_dataloader, 
+                    loss_fn=loss_fn, 
+                    optimizer=optimizer, 
+                    epochs=NUM_EPOCHS, 
+                    device=device)
+            
+    # save the model with help from utils.py
+    utils.save_model(model=model, target_dir="models", model_name="tinyfood-effnet.pt")
+
+if __name__ == '__main__':
+    main()

utils.py

@@ -0,0 +1,152 @@
+"""
+contains various utility functions for pytorch model training and saving
+"""
+import torch 
+from pathlib import Path 
+import matplotlib.pyplot as plt
+import torchvision
+from PIL import Image
+from torch.utils.tensorboard.writer import SummaryWriter
+
+def save_model(model: torch.nn.Module, 
+              target_dir: str, 
+              model_name: str):
+    """Saves a pytorch model to a target directory
+
+    Args:
+        model: target pytorch model
+        target_dir: string of target directory path to store the saved models 
+        model_name: a filename for the saved model. Should be included either ".pth" or ".pt" as 
+        the file extension.
+    """
+    # create target directory 
+    target_dir_path = Path(target_dir)
+    target_dir_path.mkdir(parents=True, exist_ok=True)
+
+    # create model save path 
+    assert model_name.endswith(".pth") or model_name.endswith(".pt"), "model name should end with .pt or .pth"
+    model_save_path = target_dir_path / model_name
+
+    # save the model state_dict()
+    print(f"[INFO] Saving model to: {model_save_path}")
+    torch.save(obj=model.state_dict(), f=model_save_path)
+
+def pred_and_plot_image(
+    model: torch.nn.Module,
+    image_path: str,
+    class_names: list[str] = None,
+    transform=None,
+    device: torch.device = "cuda" if torch.cuda.is_available() else "cpu",
+):
+    """Makes a prediction on a target image with a trained model and plots the image.
+
+    Args:
+        model (torch.nn.Module): trained PyTorch image classification model.
+        image_path (str): filepath to target image.
+        class_names (List[str], optional): different class names for target image. Defaults to None.
+        transform (_type_, optional): transform of target image. Defaults to None.
+        device (torch.device, optional): target device to compute on. Defaults to "cuda" if torch.cuda.is_available() else "cpu".
+    
+    Returns:
+        Matplotlib plot of target image and model prediction as title.
+
+    Example usage:
+        pred_and_plot_image(model=model,
+                            image="some_image.jpeg",
+                            class_names=["class_1", "class_2", "class_3"],
+                            transform=torchvision.transforms.ToTensor(),
+                            device=device)
+    """
+
+    # 1. Load in image and convert the tensor values to float32
+    img_list = Image.open(image_path)
+
+    # 2. Divide the image pixel values by 255 to get them between [0, 1]
+    # target_image = target_image / 255.0
+
+    # 3. Transform if necessary
+    if transform:
+        target_image = transform(img_list)
+
+    # 4. Make sure the model is on the target device
+    model.to(device)
+
+    # 5. Turn on model evaluation mode and inference mode
+    model.eval()
+    with torch.inference_mode():
+        # Add an extra dimension to the image
+        target_image = target_image.unsqueeze(dim=0)
+
+        # Make a prediction on image with an extra dimension and send it to the target device
+        target_image_pred = model(target_image.to(device))
+
+    # 6. Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)
+    target_image_pred_probs = torch.softmax(target_image_pred, dim=1)
+
+    # 7. Convert prediction probabilities -> prediction labels
+    target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)
+
+    # 8. Plot the image alongside the prediction and prediction probability
+    plt.imshow(
+        target_image.squeeze().permute(1, 2, 0)
+    )  # make sure it's the right size for matplotlib
+    if class_names:
+        title = f"Pred: {class_names[target_image_pred_label.cpu()]} | Prob: {target_image_pred_probs.max().cpu():.3f}"
+    else:
+        title = f"Pred: {target_image_pred_label} | Prob: {target_image_pred_probs.max().cpu():.3f}"
+    plt.title(title)
+    plt.axis(False)
+
+def set_seeds(seed: int=42):
+    """Sets random sets for torch operations.
+
+    Args:
+        seed (int, optional): Random seed to set. Defaults to 42.
+    """
+    # Set the seed for general torch operations
+    torch.manual_seed(seed)
+    # Set the seed for CUDA torch operations (ones that happen on the GPU)
+    torch.cuda.manual_seed(seed)
+
+
+def create_writer(experiment_name: str, model_name: str, extra: str=None) -> torch.utils.tensorboard.writer.SummaryWriter(): # type: ignore
+    """
+    creates a torch.utils.tensorboard.writer.SummaryWriter() instance saving to a
+    specific log_dir.
+
+    log_dir is a combination of runs/timestamp/experiment_name/model_name/extra.
+
+    where timestamp is the current date in YYYY-MM-DD format.
+
+    Args:
+        experiment_name (str): Name of experiment
+        model_name (str): model name
+        extra (str, optional): anything extra to add to the directory. Defaults is None
+
+    Returns:
+        torch.utils.tensorboard.writer.SummaryWriter(): Instance of a writer saving to log_dir
+
+    Examples usage:
+        this is gonna create writer saving to "runs/2022-06-04/data_10_percent/effnetb2/5_epochs"
+
+    writer = create_writer(experiment_name="data_10_percent", model_name="effnetb2", extra="5_epochs")
+
+    This is the same as:
+    writer = SummaryWriter(log_dir="runs/2022-06-04/data_10_percent/effnetb2/5_epochs")
+    """
+
+    from datetime import datetime
+    import os
+
+    # get the timestamp
+    timestamp = datetime.now().strftime("%Y-%m-%d")
+
+    if extra:
+        # create log directory path
+        log_dir = os.path.join("runs", timestamp, experiment_name, model_name, extra)
+    else:
+        log_dir = os.path.join("runs", timestamp, experiment_name, model_name)
+
+    print(f"[INFO] Created SummaryWriter(), saving to: {log_dir}")
+    
+    return SummaryWriter(log_dir=log_dir)