Merge pull request #69 from SFI-Visual-Intelligence/usps-inputshape

Usps inputshape, resolve #67
LGTM

Author: Seilmast

main.py

@@ -30,10 +30,12 @@ def main():
 
     device = args.device
 
-    if args.dataset.lower() in ["usps_0-6", "usps_7-9"]:
+
+    if "usps" in args.dataset.lower():
+
         transform = transforms.Compose(
             [
-                transforms.Resize((16, 16)),
+                transforms.Resize((28, 28)),
                 transforms.ToTensor(),
             ]
         )
@@ -45,6 +47,7 @@ def main():
         data_dir=args.datafolder,
         transform=transform,
         val_size=args.val_size,
+
     )
 
     train_metrics = MetricWrapper(
@@ -126,6 +129,7 @@ def main():
         project=args.run_name,
         tags=[args.modelname, args.dataset],
         config=args,
+
     )
     wandb.watch(model)
 

utils/arg_parser.py

@@ -33,6 +33,7 @@ def get_args():
         help="Whether model should be saved or not.",
     )
 
+
     # Data/Model specific values
     parser.add_argument(
         "--modelname",
@@ -82,6 +83,7 @@ def get_args():
         "--macro_averaging",
         action="store_true",
         help="If the flag is included, the metrics will be calculated using macro averaging.",
+
     )
 
     # Training specific values

utils/dataloaders/svhn.py

@@ -1,5 +1,6 @@
 import os
 
+
 import h5py
 import numpy as np
 from PIL import Image
@@ -29,6 +30,7 @@ def __init__(
             AssertionError: If the split is not 'train' or 'test'.
         """
         super().__init__()
+
         self.data_path = data_path
         self.split = "train" if train else "test"
 
@@ -55,6 +57,7 @@ def _download_data(self, path: str):
             path (str): The directory where the dataset will be downloaded.
         """
         print(f"Downloading SVHN data into {path}")
+
         SVHN(path, split=self.split, download=True)
         data = loadmat(os.path.join(path, f"{self.split}_32x32.mat"))
 
@@ -92,8 +95,8 @@ def __getitem__(self, index):
             img = Image.fromarray(h5f["images"][index])
 
         if self.nr_channels == 1:
-            img = img.convert("L")
 
+            img = img.convert("L")
         if self.transforms is not None:
             img = self.transforms(img)
 

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)