optimised

architectures/ResNet18.py

@@ -0,0 +1,22 @@
+
+import torch.nn as nn
+from torchvision import models
+
+# Base model
+from architectures.Model import Model
+
+class ResNet18(Model):
+    
+    def get(self):
+        m = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
+
+        # freez all layers
+        #for param in m.parameters():
+        #    param.requires_grad = False
+
+        # unfreez the last two
+        #for param in m.layer3.parameters(): param.requires_grad = True
+        #for param in m.layer4.parameters(): param.requires_grad = True
+        
+        m.fc = nn.Linear(m.fc.in_features, self.size)
+        return m