optimised

architectures/EfficentNet.py

@@ -0,0 +1,19 @@
+
+import torch.nn as nn
+from torchvision import models
+
+# Base model
+from architectures.Model import Model
+
+class EfficientNet(Model):
+    
+    def get(self):
+
+        m = models.efficientnet_b1(weights=models.EfficientNet_B1_Weights.DEFAULT)
+
+        # Unfreeze the last block for a lighter touch
+        for param in m.features[-1].parameters(): param.requires_grad = True
+
+        # Standard classifier fix
+        m.classifier[1] = nn.Linear(m.classifier[1].in_features, self.size)
+        return m