finds number of channels based on dataset. Adds num_classes to dataset

Author: c-salomonsen

main.py

@@ -108,7 +108,7 @@ def main():
     parser.add_argument(
         "--device",
         type=str,
-        default="cuda",
+        default="cpu",
         choices=["cuda", "cpu", "mps"],
         help="Which device to run the training on.",
     )
@@ -124,10 +124,6 @@ def main():
 
     device = args.device
 
-    # load model
-    model = load_model(args.modelname)
-    model.to(device)
-
     metrics = MetricWrapper(*args.metric)
 
     # Dataset
@@ -143,6 +139,20 @@ def main():
         data_path=args.datafolder,
     )
 
+    # Find number of channels in the dataset
+    if len(traindata[0][0].shape) == 2:
+        channels = 1
+    else:
+        channels = traindata[0][0].shape[0]
+
+    # load model
+    model = load_model(
+        args.modelname,
+        in_channels=channels,
+        num_classes=traindata.num_classes,
+    )
+    model.to(device)
+
     trainloader = DataLoader(traindata,
                              batch_size=args.batchsize,
                              shuffle=True,

utils/dataloaders/usps_0_6.py

@@ -36,6 +36,8 @@ class USPSDataset0_6(Dataset):
         A function/transform that takes in a sample and returns a transformed version.
     idx : numpy.ndarray
         Indices of samples with labels 0-6.
+    num_classes : int
+        Number of classes in the dataset
 
     Methods
     -------
@@ -71,6 +73,7 @@ def __init__(
         super().__init__()
         self.path = list(data_path.glob("*.h5"))[0]
         self.transform = transform
+        self.num_classes = 7
 
         if download:
             raise NotImplementedError("Download functionality not implemented.")