Update model.py

model.py

@@ -1,6 +1,5 @@
 import torch.nn as nn
-
-from config import IMAGE_SIZE
+from torchvision import models
 
 
 class SimpleCNN(nn.Module):
@@ -15,27 +14,17 @@ class SimpleCNN(nn.Module):
     ):
         super().__init__()
 
-        padding = kernel_size // 2
-
-        self.features = nn.Sequential(
-            nn.Conv2d(3, conv1_channels, kernel_size=kernel_size, padding=padding),
-            nn.ReLU(),
-            nn.MaxPool2d(2),
-
-            nn.Conv2d(conv1_channels, conv2_channels, kernel_size=kernel_size, padding=padding),
-            nn.ReLU(),
-            nn.MaxPool2d(2),
-        )
+        weights = models.ResNet18_Weights.DEFAULT
+        self.backbone = models.resnet18(weights=weights)
 
-        flattened_dim = conv2_channels * (IMAGE_SIZE // 4) * (IMAGE_SIZE // 4)
-
-        self.classifier = nn.Sequential(
-            nn.Flatten(),
-            nn.Linear(flattened_dim, fc_dim),
+        in_features = self.backbone.fc.in_features
+        self.backbone.fc = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(in_features, fc_dim),
             nn.ReLU(),
             nn.Dropout(dropout),
             nn.Linear(fc_dim, num_classes),
         )
 
     def forward(self, x):
-        return self.classifier(self.features(x))
+        return self.backbone(x)