optimised

architectures/DenseNet121.py

@@ -0,0 +1,15 @@
+
+import torch.nn as nn
+from torchvision import models
+from architectures.Model import Model
+
+class DenseNet121(Model):
+    def get(self):
+
+        # load pretrained
+        m = models.densenet121(weights=models.DenseNet121_Weights.DEFAULT)
+        # will modify only the final layers
+        num_ftrs = m.classifier.in_features
+        m.classifier = nn.Linear(num_ftrs, self.size)
+        
+        return m