Update model.py

model.py

@@ -2,15 +2,13 @@ import torch.nn as nn
 from torchvision import models
 
 
-class SimpleCNN(nn.Module):
+class ResNet18Classifier(nn.Module):
     def __init__(
         self,
         num_classes: int,
-        conv1_channels: int = 16,
-        conv2_channels: int = 32,
-        kernel_size: int = 3,
-        dropout: float = 0.2,
-        fc_dim: int = 128,
+        dropout: float = 0.5,
+        fc_dim: int = 256,
+        freeze_backbone: bool = True,
     ):
         super().__init__()
 
@@ -18,6 +16,11 @@ class SimpleCNN(nn.Module):
         self.backbone = models.resnet18(weights=weights)
 
         in_features = self.backbone.fc.in_features
+
+        if freeze_backbone:
+            for param in self.backbone.parameters():
+                param.requires_grad = False
+
         self.backbone.fc = nn.Sequential(
             nn.Dropout(dropout),
             nn.Linear(in_features, fc_dim),
@@ -26,5 +29,8 @@ class SimpleCNN(nn.Module):
             nn.Linear(fc_dim, num_classes),
         )
 
+        for param in self.backbone.fc.parameters():
+            param.requires_grad = True
+
     def forward(self, x):
         return self.backbone(x)