added loading and testing the model again

architectures/GoogleNet.py

@@ -9,6 +9,23 @@ class GoogleNet(Model):
     
     def get(self):
 
-        m = models.googlenet(weights =  models.GoogLeNet_Weights.DEFAULT)
+        m = models.googlenet(weights=models.GoogLeNet_Weights.DEFAULT)
+
+        # 1. Handle the two Auxiliary Classifiers
+        # GoogLeNet has aux1 and aux2 to help training converge
+        #if m.aux_logits:
+            #m.aux1.fc = nn.Linear(m.aux1.fc.in_features, self.size)
+            #m.aux2.fc = nn.Linear(m.aux2.fc.in_features, self.size)
+
+        # 2. Handle the Main Classifier
         m.fc = nn.Linear(m.fc.in_features, self.size)
+        
+        #for param in m.parameters():
+        #    param.requires_grad = False
+            
+        # Unfreezing the final stages for identity recognition
+        #for name, param in m.named_parameters():
+        #    if "inception5" in name or "fc" in name:
+        #        param.requires_grad = True
+
         return m

architectures/Model.py

@@ -62,9 +62,7 @@ class Model(ABC):
 
 
     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)
 
@@ -72,8 +70,29 @@ class Model(ABC):
         if filename is None:
             filename = f"{self.__class__.__name__.lower()}.pth"
         
+        if not filename.endswith('.pth'):
+            filename += '.pth'
+
         save_path = save_dir / filename
         torch.save(self.model.state_dict(), save_path)
+        print(f'Model saved to {save_path}')
+
+
+
+    def load(self, arch):
+ 
+        file_path = Path("trained_models") / f'{arch.name.lower()}.pth'
+        # does file exist
+        if not file_path.exists():
+            raise FileNotFoundError(f'No checkpoint found at: {file_path}')
+            
+        # Load the weights
+        state_dict = torch.load(file_path, map_location=self.device, weights_only=True)
+        
+        self.model.load_state_dict(state_dict)
+        self.model.to(self.device)
+        print(f'Model loaded from {file_path}')
+
 
 
     # Using the factory patern here