From cfa4da697eb7f8ceb969e0a8021224049c9f6a5d Mon Sep 17 00:00:00 2001
From: Tinsae <tinsaekahsay@gmail.com>
Date: Tue, 5 May 2026 22:42:27 +0200
Subject: [PATCH] added loading and testing the model again

---
 ReadME.md                  |  6 +++++-
 Tune.py                    | 34 +++++++++++++++++++++-------------
 architectures/GoogleNet.py | 19 ++++++++++++++++++-
 architectures/Model.py     | 25 ++++++++++++++++++++++---
 4 files changed, 66 insertions(+), 18 deletions(-)

diff --git a/ReadME.md b/ReadME.md
index f44ec5d..080db7d 100644
--- a/ReadME.md
+++ b/ReadME.md
@@ -36,11 +36,15 @@ Root_dir/
 
 ```
 
-once this is manually done, We can run finetunning a selected model. For now, 4 models are implemented.
+once this is manually done, We can run finetunning a selected model. For now, 8 models are implemented.
 - ResNet-18
 - ResNet-50
 - DenseNet121
 - Inception
+- GoogleNet
+- ShuffleNet
+- EfficientNet
+- WideResNet
 
 ## Fine tuning
 ### Preparation
diff --git a/Tune.py b/Tune.py
index 402e658..8cc66e4 100644
--- a/Tune.py
+++ b/Tune.py
@@ -8,7 +8,7 @@ from IdentitySubset import IdentitySubset
 from architectures.Model import Model, Architecture
 
 # numbre of classes 
-CLASS_SIZE = 50
+CLASS_SIZE = 20
 # batch 
 BATCH_SIZE = 16
 
@@ -21,7 +21,7 @@ TRAINING_SMPLE = 28
 
 # learning rate
 LR_RATE = 0.0001
-EPOCHS = 20
+EPOCHS = 16
 
 # depends on model architecture
 # ResNet, DenseNet = 224
@@ -36,7 +36,7 @@ RESOLUTION = 224
 #  - GOOGLENET
 #  - EFFICIENTNET
 #  - SHUFFLENET
-arch = Architecture.GOOGLENET
+arch = Architecture.RESNET50
 
 # DATA PREPARATION
 # load data set and prepare
@@ -88,19 +88,15 @@ model = Model.create(
     device = device, 
     size = CLASS_SIZE)
 
-# FINETUNING
+# we may need to load existing model or finetune
 model.train(
-    epochs = EPOCHS, 
-    loader = train_loader, 
-    rate = LR_RATE)
+        epochs = EPOCHS, 
+        loader = train_loader, 
+        rate = LR_RATE)
 
-# save.
+    # save.
 model.save(filename=arch.name.lower())
-'''
-torch.save(
-    model.get().state_dict(), 
-    f'trained/{arch.name.lower()}.pth'
-)'''
+
 
 # done tuning
 print('Model saved!')
@@ -125,3 +121,15 @@ print(f"Total test images for these {CLASS_SIZE} classes: {len(test_data)}")
 # Evaluate
 model.evaluate(
     loader = test_loader)
+
+# test again
+reloaded = Model.create(
+    arch=arch,
+    device = device,
+    size = CLASS_SIZE
+    )
+reloaded.load(arch = arch)
+print("Evaluating loaded")
+reloaded.evaluate(
+    loader = test_loader
+)
diff --git a/architectures/GoogleNet.py b/architectures/GoogleNet.py
index 33a535e..8dd9c9c 100644
--- a/architectures/GoogleNet.py
+++ b/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
\ No newline at end of file
diff --git a/architectures/Model.py b/architectures/Model.py
index 1c20997..50d506d 100644
--- a/architectures/Model.py
+++ b/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