added more models for testing

architectures/Model.py

@@ -6,17 +6,18 @@ import torch.optim as optim
 import time
 import numpy as np
 from sklearn.metrics import classification_report
+from pathlib import Path
 
 class Model(ABC):
     def __init__(self, device, size):
         self.device = device
         self.size = size
-        self.model = self.get()
+        self.model = self.get().to(self.device)
 
     @abstractmethod
     def get(self):
         # return the model
-        return self.model
+        pass
 
     def train(self, epochs, loader, rate):
         criterion = nn.CrossEntropyLoss()
@@ -60,6 +61,21 @@ class Model(ABC):
         print(classification_report(all_labels, all_preds, zero_division=0))
 
 
+    def save(self, filename=None):
+        """
+        Saves the model state_dict. Creates the directory if it doesn't exist.
+        """
+        save_dir = Path("trained_models")
+        save_dir.mkdir(parents=True, exist_ok=True)
+
+        # 2. Determine filename (Default to class name if not provided)
+        if filename is None:
+            filename = f"{self.__class__.__name__.lower()}.pth"
+        
+        save_path = save_dir / filename
+        torch.save(self.model.state_dict(), save_path)
+
+
     # Using the factory patern here
     @staticmethod
     def create(arch, device, size):
@@ -74,8 +90,8 @@ class Model(ABC):
             
             # ResNet50
             case Architecture.RESNET50:
-                from architectures.ResNet18 import ResNet18
-                return ResNet18(device, size)
+                from architectures.ResNet50 import ResNet50
+                return ResNet50(device, size)
             
             # INCEPTION
             case Architecture.INCEPTION:
@@ -94,6 +110,14 @@ class Model(ABC):
             case Architecture.EFFICIENTNET:
                 from architectures.EfficentNet import EfficientNet
                 return EfficientNet(device, size)
+            #ShuffleNet
+            case Architecture.SHUFFLENET:
+                from architectures.ShuffleNet import ShuffleNet
+                return ShuffleNet(device, size)
+            # wide ResNet
+            case Architecture.WIDE_RESNET:
+                from architectures.WideResNet import WideResNet
+                return WideResNet(device, size)
             case _:
                 raise ValueError(f"Unknown model: {arch}")
 
@@ -106,4 +130,6 @@ class Architecture(Enum):
     INCEPTION = auto()
     DENSENET121 = auto()
     GOOGLENET = auto()
-    EFFICIENTNET = auto()
+    EFFICIENTNET = auto()
+    SHUFFLENET = auto()
+    WIDE_RESNET = auto()