optimised

architectures/ShuffleNet.py

@@ -0,0 +1,17 @@
+
+
+
+import torch.nn as nn
+from torchvision import models
+
+# Base model
+from architectures.Model import Model
+
+class ShuffleNet(Model):
+    
+    def get(self):
+        m = models.shufflenet_v2_x1_0(weights=models.ShuffleNet_V2_X1_0_Weights.DEFAULT)
+        
+        num_ftrs = m.fc.in_features
+        m.fc = nn.Linear(num_ftrs, self.size)
+        return m