optimised

architectures/GoogleNet.py

@@ -0,0 +1,31 @@
+
+import torch.nn as nn
+from torchvision import models
+
+# Base model
+from architectures.Model import Model
+
+class GoogleNet(Model):
+    
+    def get(self):
+
+        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