model fixed

model.py

@@ -1,6 +1,7 @@
 import pytorch_lightning as pl
 import torch
 import torchmetrics
+import torchvision.models as models
 from simple_parsing import ArgumentParser
 from torch import nn
 from torch.nn import functional as F
@@ -14,46 +15,48 @@ args = args_namespace.options
 
 # Model class
 class Model(nn.Module):
-    def __init__(self):
+    def __init__(self, input_shape, weights=args.weights):
         super().__init__()
 
-        self.conv1 = nn.Conv2d(3, 32, 5)
-        self.conv2 = nn.Conv2d(32, 64, 5)
-        self.conv3 = nn.Conv2d(64, 128, 3)
-        self.dropout1 = nn.Dropout2d(0.25)
-        self.dropout2 = nn.Dropout2d(0.5)
+        self.feature_extractor = models.resnet18(weights=weights)
+
+        if weights:
+            # layers are frozen by using eval()
+            self.feature_extractor.eval()
+            # freeze params
+            for param in self.feature_extractor.parameters():
+                param.requires_grad = False
+
+        n_size = self._get_conv_output(input_shape)
+
+        self.classifier = nn.Linear(n_size, args.num_classes)
 
-        x = torch.randn(3, 224, 224).view(-1, 3, 224, 224)
-        self._to_linear = None
-        self.convs(x)
+    # returns the size of the output tensor going into the Linear layer from the conv block.
+    def _get_conv_output(self, shape):
+        batch_size = 1
+        tmp_input = torch.autograd.Variable(torch.rand(batch_size, *shape))
 
-        self.fc1 = nn.Linear(self._to_linear, 128)
-        self.fc2 = nn.Linear(128, args.num_classes)
+        output_feat = self.convs(tmp_input)
+        n_size = output_feat.data.view(batch_size, -1).size(1)
+        return n_size
 
     def convs(self, x):
-        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
-        x = self.dropout1(x)
-        x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2))
-        x = self.dropout2(x)
-        x = F.max_pool2d(F.relu(self.conv3(x)), (2, 2))
-
-        if self._to_linear is None:
-            self._to_linear = x[0].shape[0] * x[0].shape[1] * x[0].shape[2]
+        x = self.feature_extractor(x)
         return x
 
     def forward(self, x):
+
         x = self.convs(x)
-        x = x.view(-1, self._to_linear)
-        x = F.relu(self.fc1(x))
-        x = self.fc2(x)
-        return F.log_softmax(x, dim=1)
+        x = x.view(x.size(0), -1)
+        x = self.classifier(x)
+        return x
 
 
 class Classifier(pl.LightningModule):
     def __init__(self):
         super().__init__()
 
-        self.model = Model()
+        self.model = Model(input_shape=args.input_shape)
         self.accuracy = torchmetrics.Accuracy(
             task="multiclass", num_classes=args.num_classes
         )
@@ -62,13 +65,13 @@ class Classifier(pl.LightningModule):
         x = self.model(x)
         return x
 
-    def nll_loss(self, logits, labels):
-        return F.nll_loss(logits, labels)
+    def ce_loss(self, logits, labels):
+        return F.cross_entropy(logits, labels)
 
     def training_step(self, train_batch, batch_idx):
         x, y = train_batch
         logits = self.model(x)
-        loss = self.nll_loss(logits, y)
+        loss = self.ce_loss(logits, y)
         acc = self.accuracy(logits, y)
         self.log("accuracy/train_accuracy", acc)
         self.log("loss/train_loss", loss)
@@ -77,11 +80,18 @@ class Classifier(pl.LightningModule):
     def validation_step(self, val_batch, batch_idx):
         x, y = val_batch
         logits = self.model(x)
-        loss = self.nll_loss(logits, y)
+        loss = self.ce_loss(logits, y)
         acc = self.accuracy(logits, y)
         self.log("accuracy/val_accuracy", acc)
         self.log("loss/val_loss", loss)
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(self.parameters(), lr=args.learning_rate)
-        return optimizer
+        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer, mode="min", patience=7
+        )
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": scheduler,
+            "monitor": "loss/val_loss",
+        }

models/checkpoint.ckpt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:85f9ff02a03ded56ff20903f0227f017ec351a9946b7fa0a1b7c33b7107427d6
-size 124442154
+oid sha256:293e1c26f4f6ec48d58824658c41e3a4509d1313b2bdf15177020893e0ed1df5
+size 140535623