Update

configs.py

@@ -3,14 +3,16 @@ import torch
 from torchvision import transforms
 from torch.utils.data import Dataset
 from models import *
-
+import torch.nn as nn
+from torchvision.models import squeezenet1_0, SqueezeNet1_0_Weights
+from torchvision.models import squeezenet1_0
 # Constants
 RANDOM_SEED = 123
 BATCH_SIZE = 32
-NUM_EPOCHS = 100
+NUM_EPOCHS = 40
 LEARNING_RATE = 0.00017588413773574044
 STEP_SIZE = 10
-GAMMA = 0.3
+GAMMA = 0.9
 DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 NUM_PRINT = 100
 TASK = 1
@@ -22,10 +24,31 @@ NUM_CLASSES = 7
 EARLY_STOPPING_PATIENCE = 20
 CLASSES = ['Alzheimer Disease', 'Cerebral Palsy', 'Dystonia', 'Essential Tremor', 'Healthy', 'Huntington Disease', 'Parkinson Disease']
 MODEL_SAVE_PATH = "output/checkpoints/model.pth"
-MODEL = squeezenet1_0(num_classes=NUM_CLASSES)
 
-print(CLASSES)
 
+    
+class SqueezeNet1_0WithDropout(nn.Module):
+    def __init__(self, num_classes=1000):
+        super(SqueezeNet1_0WithDropout, self).__init__()
+        squeezenet = squeezenet1_0(weights=SqueezeNet1_0_Weights) 
+        self.features = squeezenet.features
+        self.classifier = nn.Sequential(
+            nn.Conv2d(512, num_classes, kernel_size=1),
+            nn.BatchNorm2d(num_classes),  # add batch normalization
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d((1, 1))
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.classifier(x)
+        x = torch.flatten(x, 1)
+        return x
+    
+    
+    
+MODEL = SqueezeNet1_0WithDropout(num_classes=7)
+print(CLASSES)
 
 preprocess = transforms.Compose(
     [
@@ -47,4 +70,4 @@ class CustomDataset(Dataset):
 
     def __getitem__(self, idx):
         img, label = self.data[idx]
-        return img, label
+        return img, label

convert.py

@@ -0,0 +1,7 @@
+import torch
+import onnx2tf
+from configs import *
+
+torch.onnx.export(model=MODEL, args=torch.randn(1, 3, 64, 64), f='output/checkpoints/model.onnx', verbose=True, input_names=['input'], output_names=['output'])
+
+onnx2tf.convert(input_onnx_file_path='output/checkpoints/model.onnx', output_folder_path='output/checkpoints/converted/')

train.py

@@ -153,6 +153,10 @@ def main_training_loop():
             break
 
     # Save the model
+    MODEL_SAVE_PATH = "output/checkpoints/model.pth"
+
+    # Ensure the parent directory exists
+    os.makedirs(os.path.dirname(MODEL_SAVE_PATH), exist_ok=True)
     torch.save(model.state_dict(), MODEL_SAVE_PATH)
     print("Model saved at", MODEL_SAVE_PATH)
 
@@ -190,4 +194,4 @@ def main_training_loop():
 
 
 if __name__ == "__main__":
-    main_training_loop()
+    main_training_loop()

tuning.py

@@ -12,7 +12,7 @@ from torch.utils.tensorboard import SummaryWriter
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 EPOCHS = 10
 N_TRIALS = 50
-TIMEOUT = None
+TIMEOUT = 900
 
 # Create a TensorBoard writer
 writer = SummaryWriter(log_dir="output/tensorboard/tuning")