Modify christianmodel to be input shape agnostic, default img shape usps: 28x28

Author: c-salomonsen

main.py

@@ -1,11 +1,11 @@
 import numpy as np
 import torch as th
 import torch.nn as nn
+import wandb
 from torch.utils.data import DataLoader
 from torchvision import transforms
 from tqdm import tqdm
 
-import wandb
 from utils import MetricWrapper, createfolders, get_args, load_data, load_model
 
 
@@ -29,30 +29,30 @@ def main():
 
     device = args.device
 
-    if args.dataset.lower() in ["usps_0-6", "uspsh5_7_9"]:
-        augmentations = transforms.Compose(
+    if "usps" in args.dataset.lower():
+        transform = transforms.Compose(
             [
-                transforms.Resize((16, 16)),
+                transforms.Resize((28, 28)),
                 transforms.ToTensor(),
             ]
         )
     else:
-        augmentations = transforms.Compose([transforms.ToTensor()])
+        transform = transforms.Compose([transforms.ToTensor()])
 
     # Dataset
     traindata = load_data(
         args.dataset,
         train=True,
         data_path=args.datafolder,
         download=args.download_data,
-        transform=augmentations,
+        transform=transform,
     )
     validata = load_data(
         args.dataset,
         train=False,
         data_path=args.datafolder,
         download=args.download_data,
-        transform=augmentations,
+        transform=transform,
     )
 
     metrics = MetricWrapper(*args.metric, num_classes=traindata.num_classes)
@@ -113,11 +113,11 @@ def main():
 
     # wandb.login(key=WANDB_API)
     wandb.init(
-            entity="ColabCode-org",
-            # entity="FYS-8805 Exam",
-            project="Test", 
-            tags=[args.modelname, args.dataset]
-            )
+        entity="ColabCode-org",
+        # entity="FYS-8805 Exam",
+        project="Test",
+        tags=[args.modelname, args.dataset],
+    )
     wandb.watch(model)
     exit()
     for epoch in range(args.epoch):

utils/models/christian_model.py

@@ -3,6 +3,18 @@
 
 
 class CNNBlock(nn.Module):
+    """
+    CNN block with Conv2d, MaxPool2d, and ReLU.
+
+    Args
+    ----
+
+    in_channels : int
+        Number of input channels.
+    out_channels : int
+        Number of output channels.
+    """
+
     def __init__(self, in_channels, out_channels):
         super().__init__()
 
@@ -22,6 +34,37 @@ def forward(self, x):
         return x
 
 
+def find_fc_input_shape(image_shape, *cnn_layers):
+    """
+    Find the shape of the input to the fully connected layer.
+
+    Code inspired by @Seilmast (https://github.com/SFI-Visual-Intelligence/Collaborative-Coding-Exam/issues/67#issuecomment-2651212254)
+
+    Args
+    ----
+    image_shape : tuple(int, int, int)
+        Shape of the input image (C, H, W).
+    cnn_layers : nn.Module
+        List of CNN layers.
+
+    Returns
+    -------
+    int
+        Number of elements in the input to the fully connected layer.
+    """
+
+    dummy_img = torch.randn(1, *image_shape)
+    with torch.no_grad():
+        x = cnn_layers[0](dummy_img)
+
+        for layer in cnn_layers[1:]:
+            x = layer(x)
+
+        x = x.view(x.size(0), -1)
+
+    return x.size(1)
+
+
 class ChristianModel(nn.Module):
     """Simple CNN model for image classification.
 
@@ -57,7 +100,9 @@ def __init__(self, image_shape, num_classes):
         self.cnn1 = CNNBlock(C, 50)
         self.cnn2 = CNNBlock(50, 100)
 
-        self.fc1 = nn.Linear(100 * 4 * 4, num_classes)
+        fc_input_shape = find_fc_input_shape(image_shape, self.cnn1, self.cnn2)
+
+        self.fc1 = nn.Linear(fc_input_shape, num_classes)
 
     def forward(self, x):
         x = self.cnn1(x)
@@ -70,9 +115,10 @@ def forward(self, x):
 
 
 if __name__ == "__main__":
-    model = ChristianModel(3, 7)
+    x = torch.randn(3, 3, 28, 28)
+
+    model = ChristianModel(x.shape[1:], 7)
 
-    x = torch.randn(3, 3, 16, 16)
     y = model(x)
 
     print(y)