Used ruff and isort to fix code structure

Author: Seilmast

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
 
 
@@ -22,13 +22,13 @@ def main():
     ------
 
     """
-    
+
     args = get_args()
-    
+
     createfolders(args.datafolder, args.resultfolder, args.modelfolder)
-    
+
     device = args.device
-    
+
     if args.dataset.lower() in ["usps_0-6", "uspsh5_7_9"]:
         augmentations = transforms.Compose(
             [
@@ -38,7 +38,7 @@ def main():
         )
     else:
         augmentations = transforms.Compose([transforms.ToTensor()])
-    
+
     # Dataset
     traindata = load_data(
         args.dataset,
@@ -54,22 +54,22 @@ def main():
         download=args.download_data,
         transform=augmentations,
     )
-    
+
     metrics = MetricWrapper(traindata.num_classes, *args.metric)
-    
+
     # Find the shape of the data, if is 2D, add a channel dimension
     data_shape = traindata[0][0].shape
     if len(data_shape) == 2:
         data_shape = (1, *data_shape)
-    
+
     # load model
     model = load_model(
         args.modelname,
         image_shape=data_shape,
         num_classes=traindata.num_classes,
     )
     model.to(device)
-    
+
     trainloader = DataLoader(
         traindata,
         batch_size=args.batchsize,
@@ -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):

tests/test_dataloaders.py

@@ -1,4 +1,5 @@
-from utils.dataloaders import USPSDataset0_6, SVHNDataset
+from utils.dataloaders import SVHNDataset, USPSDataset0_6
+
 
 def test_uspsdataset0_6():
     from pathlib import Path
@@ -32,28 +33,23 @@ def test_uspsdataset0_6():
         assert data.shape == (1, 16, 16)
         assert all(target == np.array([0, 0, 0, 0, 0, 0, 1]))
 
-        
+
 def test_svhn_dataset():
     import os
     from tempfile import TemporaryDirectory
+
     from torchvision import transforms
-    
-    with TemporaryDirectory() as tempdir: 
-        
-        trans = transforms.Compose([
-            transforms.Resize((28,28)),
-            transforms.ToTensor()
-        ])
-        
-        dataset = SVHNDataset(tempdir,
-                              train=True,
-                              transform=trans,
-                              download=True,
-                              nr_channels=1)
-        
+
+    with TemporaryDirectory() as tempdir:
+        trans = transforms.Compose([transforms.Resize((28, 28)), transforms.ToTensor()])
+
+        dataset = SVHNDataset(
+            tempdir, train=True, transform=trans, download=True, nr_channels=1
+        )
+
         assert dataset.__len__() != 0
-        assert os.path.exists(os.path.join(tempdir, 'train_32x32.mat'))
-        
+        assert os.path.exists(os.path.join(tempdir, "train_32x32.mat"))
+
         img, label = dataset.__getitem__(0)
-        assert len(img.size()) == 3 and img.size() == (1,28,28) and img.size(0) == 1
+        assert len(img.size()) == 3 and img.size() == (1, 28, 28) and img.size(0) == 1
         assert len(label.size()) == 1

tests/test_metrics.py

@@ -1,4 +1,4 @@
-from utils.metrics import Accuracy, F1Score, Precision, Recall, EntropyPrediction
+from utils.metrics import Accuracy, EntropyPrediction, F1Score, Precision, Recall
 
 
 def test_recall():
@@ -98,20 +98,25 @@ def test_accuracy():
         f"Accuracy Score: {accuracy_score.item()}"
     )
 
+
 def test_entropypred():
-    import torch as th 
-    
-    metric = EntropyPrediction(averages='mean')
+    import torch as th
+
+    metric = EntropyPrediction(averages="mean")
 
-    true_lab = th.Tensor([0,1,1,2,4,3]).reshape(6,1).type(th.LongTensor)
+    true_lab = th.Tensor([0, 1, 1, 2, 4, 3]).reshape(6, 1).type(th.LongTensor)
     pred_logits = th.nn.functional.one_hot(true_lab, 5)
-    
-    #Test for log(0) errors and expected output
+
+    # Test for log(0) errors and expected output
     assert th.abs((th.sum(metric(true_lab, pred_logits)) - 0.0)) < 1e-5
-    
-    pred_logits = th.rand(6,5)
-    metric2 = EntropyPrediction(averages='sum')
-    
-    #Test for averaging metric consistency
-    assert th.abs(th.sum(6*metric(true_lab, pred_logits) - metric2(true_lab, pred_logits))) < 1e-5
-    
+
+    pred_logits = th.rand(6, 5)
+    metric2 = EntropyPrediction(averages="sum")
+
+    # Test for averaging metric consistency
+    assert (
+        th.abs(
+            th.sum(6 * metric(true_lab, pred_logits) - metric2(true_lab, pred_logits))
+        )
+        < 1e-5
+    )

tests/test_models.py

@@ -34,19 +34,15 @@ def test_jan_model(image_shape, num_classes):
     assert y.shape == (n, num_classes), f"Shape: {y.shape}"
 
 
-@pytest.mark.parameterize(
-    "image_shape",
-    [(3,28,28)]
-)
+@pytest.mark.parameterize("image_shape", [(3, 28, 28)])
 def test_magnus_model(image_shape):
-    import torch as th 
-    
+    import torch as th
+
     n, c, h, w = 5, *image_shape
-    model = MagnusModel([h,w], 10, c)
-    
+    model = MagnusModel([h, w], 10, c)
+
     x = th.rand((n, c, h, w))
     with th.no_grad():
         y = model(x)
-        
+
     assert y.shape == (n, 10), f"Shape: {y.shape}"
-    

utils/arg_parser.py

@@ -69,20 +69,15 @@ def get_args():
         nargs="+",
         help="Which metric to use for evaluation",
     )
-    
-    parser.add_argument(
-        '--imagesize',
-        type=int,
-        default=28,
-        help='Imagesize'
-    )
-    
+
+    parser.add_argument("--imagesize", type=int, default=28, help="Imagesize")
+
     parser.add_argument(
-        '--nr_channels',
+        "--nr_channels",
         type=int,
         default=1,
-        choices=[1,3],
-        help='Number of image channels'
+        choices=[1, 3],
+        help="Number of image channels",
     )
 
     # Training specific values

utils/dataloaders/__init__.py

@@ -1,6 +1,11 @@
-__all__ = ["USPSDataset0_6", "USPSH5_Digit_7_9_Dataset", "MNISTDataset0_3", "SVHNDataset"]
+__all__ = [
+    "USPSDataset0_6",
+    "USPSH5_Digit_7_9_Dataset",
+    "MNISTDataset0_3",
+    "SVHNDataset",
+]
 
 from .mnist_0_3 import MNISTDataset0_3
+from .svhn import SVHNDataset
 from .usps_0_6 import USPSDataset0_6
 from .uspsh5_7_9 import USPSH5_Digit_7_9_Dataset
-from .svhn import SVHNDataset

utils/dataloaders/svhn.py

@@ -1,4 +1,5 @@
 import os
+
 import numpy as np
 from scipy.io import loadmat
 from torch.utils.data import Dataset
@@ -7,13 +8,13 @@
 
 class SVHNDataset(Dataset):
     def __init__(
-        self, 
-        data_path: str, 
+        self,
+        data_path: str,
         train: bool,
-        transform=None, 
-        download:bool=True, 
-        nr_channels=3
-        ):
+        transform=None,
+        download: bool = True,
+        nr_channels=3,
+    ):
         """
         Initializes the SVHNDataset object.
         Args:
@@ -26,8 +27,8 @@ def __init__(
         """
         super().__init__()
         # assert split == "train" or split == "test"
-        self.split = 'train' if train else 'test'
-    
+        self.split = "train" if train else "test"
+
         if download:
             self._download_data(data_path)
 
@@ -37,7 +38,7 @@ def __init__(
         self.images = data["X"].transpose(3, 1, 0, 2)
         self.labels = data["y"].flatten()
         self.labels[self.labels == 10] = 0
-        
+
         self.nr_channels = nr_channels
         self.transforms = transform
         self.num_classes = len(np.unique(self.labels))
@@ -50,7 +51,7 @@ def _download_data(self, path: str):
             split (str): The dataset split to download, either 'train' or 'test'.
         """
         print(f"Downloading SVHN data into {path}")
-        
+
         SVHN(path, split=self.split, download=True)
 
     def __len__(self):
@@ -73,7 +74,7 @@ def __getitem__(self, index):
 
         if self.nr_channels == 1:
             img = np.mean(img, axis=2, keepdims=True)
-        
+
         if self.transforms is not None:
             img = self.transforms(img)
 

utils/load_data.py

@@ -1,6 +1,11 @@
 from torch.utils.data import Dataset
 
-from .dataloaders import MNISTDataset0_3, USPSDataset0_6, USPSH5_Digit_7_9_Dataset, SVHNDataset
+from .dataloaders import (
+    MNISTDataset0_3,
+    SVHNDataset,
+    USPSDataset0_6,
+    USPSH5_Digit_7_9_Dataset,
+)
 
 
 def load_data(dataset: str, *args, **kwargs) -> Dataset:

utils/load_metric.py

@@ -1,6 +1,8 @@
 import copy
+
 import numpy as np
 import torch.nn as nn
+
 from .metrics import Accuracy, EntropyPrediction, F1Score, Precision, Recall
 
 
@@ -70,7 +72,7 @@ def _get_metric(self, key):
 
         match key.lower():
             case "entropy":
-                #Not dependent on knowing the number of classes
+                # Not dependent on knowing the number of classes
                 return EntropyPrediction()
             case "f1":
                 return F1Score(num_classes=self.num_classes)