Update

configs.py

@@ -6,10 +6,10 @@ from models import *
 
 # Constants
 RANDOM_SEED = 123
-BATCH_SIZE = 128
+BATCH_SIZE = 64
 NUM_EPOCHS = 100
-LEARNING_RATE = 0.04279442975996121
-OPTIMIZER_NAME = "LBFGS"
+LEARNING_RATE = 1.6317268278715415e-05
+OPTIMIZER_NAME = "Adam"
 STEP_SIZE = 10
 GAMMA = 0.5
 DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -22,7 +22,7 @@ NUM_CLASSES = 7
 # Define classes as listdir of augmented data
 CLASSES = os.listdir("data/train/augmented/Task 1/")
 MODEL_SAVE_PATH = "output/checkpoints/model.pth"
-MODEL = mobilenet_v2(num_classes=NUM_CLASSES)
+MODEL = googlenet(num_classes=NUM_CLASSES)
 
 print(CLASSES)
 

train.py

@@ -2,15 +2,11 @@ import os
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torchvision.transforms import transforms
-from torch.utils.data import DataLoader
-from torchvision.utils import make_grid
 import matplotlib.pyplot as plt
 from models import *
 from torch.utils.tensorboard import SummaryWriter
 from configs import *
 import data_loader
-import numpy as np
 
 # Set up TensorBoard writer
 writer = SummaryWriter(log_dir="output/tensorboard/training")
@@ -54,6 +50,7 @@ VAL_ACC_HIST = []
 # Training loop
 for epoch in range(NUM_EPOCHS):
     print(f"[Epoch: {epoch + 1}]")
+    print("Learning rate:", scheduler.get_last_lr()[0])
     MODEL.train()  # Set model to training mode
     running_loss = 0.0
     total_train = 0
@@ -62,7 +59,10 @@ for epoch in range(NUM_EPOCHS):
     for i, (inputs, labels) in enumerate(train_loader, 0):
         inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
         optimizer.zero_grad()
-        outputs = MODEL(inputs)
+        if MODEL.__class__.__name__ == "GoogLeNet": # the shit GoogLeNet has a different output
+            outputs = MODEL(inputs).logits
+        else:
+            outputs = MODEL(inputs)
         loss = criterion(outputs, labels)
         loss.backward()
         if OPTIMIZER_NAME == "LBFGS":

tuning.py

@@ -7,12 +7,15 @@ import torch.optim as optim
 import torch.utils.data
 from configs import *
 import data_loader
+from torch.utils.tensorboard import SummaryWriter
 
 optuna.logging.set_verbosity(optuna.logging.DEBUG)
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 EPOCHS = 10
 
+# Create a TensorBoard writer
+writer = SummaryWriter(log_dir="output/tensorboard/tuning/", )
 
 def create_data_loaders(batch_size):
     # Create or modify data loaders with the specified batch size
@@ -21,7 +24,6 @@ def create_data_loaders(batch_size):
     )
     return train_loader, valid_loader
 
-
 def objective(trial, model=MODEL):
     # Generate the model.
     model = model.to(DEVICE)
@@ -34,7 +36,7 @@ def objective(trial, model=MODEL):
 
     # Generate the optimizer.
     optimizer_name = trial.suggest_categorical("optimizer", ["Adam", "SGD"])
-    lr = trial.suggest_float("lr", 1e-5, 1e-1, log=True)
+    lr = trial.suggest_float("lr", 1e-5, 1e-3, log=True)
     optimizer = getattr(optim, optimizer_name)(model.parameters(), lr=lr)
     criterion = nn.CrossEntropyLoss()
 
@@ -45,7 +47,10 @@ def objective(trial, model=MODEL):
         for batch_idx, (data, target) in enumerate(train_loader, 0):
             data, target = data.to(DEVICE), target.to(DEVICE)
             optimizer.zero_grad()
-            output = model(data)
+            if model.__class__.__name__ == "GoogLeNet": # the shit GoogLeNet has a different output
+                output = model(data).logits
+            else:
+                output = model(data)
             loss = criterion(output, target)
             loss.backward()
             if optimizer_name == "LBFGS":
@@ -66,6 +71,19 @@ def objective(trial, model=MODEL):
 
         accuracy = correct / len(valid_loader.dataset)
 
+        # Log hyperparameters and accuracy to TensorBoard
+        writer.add_scalar("Accuracy", accuracy, trial.number)
+        writer.add_hparams(
+            {
+                "batch_size": batch_size,
+                "optimizer": optimizer_name,
+                "lr": lr
+            },
+            {
+                "accuracy": accuracy
+            }
+        )
+
         # Print hyperparameters and accuracy
         print("Hyperparameters: ", trial.params)
         print("Accuracy: ", accuracy)
@@ -77,7 +95,6 @@ def objective(trial, model=MODEL):
 
     return accuracy
 
-
 if __name__ == "__main__":
     pruner = optuna.pruners.HyperbandPruner()
     study = optuna.create_study(direction="maximize", pruner=pruner, study_name="handetect")
@@ -99,3 +116,6 @@ if __name__ == "__main__":
     print("  Params: ")
     for key, value in trial.params.items():
         print("    {}: {}".format(key, value))
+
+    # Close TensorBoard writer
+    writer.close()