Updated docstrings

Author: Seilmast

utils/dataloaders/svhn.py

@@ -16,27 +16,28 @@ def __init__(
         nr_channels=3,
     ):
         """
-        Initializes the SVHNDataset object.
+        Initializes the SVHNDataset object for loading the Street View House Numbers (SVHN) dataset.
         Args:
-            data_path (str): Path to where the data lies. If download_data is set to True, this is where the data will be downloaded.
-            transforms: Torch composite of transformations which are to be applied to the dataset images.
-            download_data (bool): If True, downloads the dataset to the specified data_path.
-            split (str): The dataset split to use, either 'train' or 'test'.
+            data_path (str): Path to where the data is stored. If `download` is set to True, this is where the data will be downloaded.
+            train (bool): If True, loads the training split of the dataset; otherwise, loads the test split.
+            transform (callable, optional): A function/transform to apply to the images.
+            download (bool): If True, downloads the dataset to the specified `data_path`.
+            nr_channels (int): Number of channels in the images. Default is 3 for RGB images.
         Raises:
             AssertionError: If the split is not 'train' or 'test'.
         """
         super().__init__()
-        # assert split == "train" or split == "test"
         self.split = "train" if train else "test"
 
         if download:
             self._download_data(data_path)
 
         data = loadmat(os.path.join(data_path, f"{self.split}_32x32.mat"))
 
-        # Images on the form N x H x W x C
+        # Reform images to the form N x H x W x C
         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
@@ -45,13 +46,11 @@ def __init__(
 
     def _download_data(self, path: str):
         """
-        Downloads the SVHN dataset.
+        Downloads the SVHN dataset to the specified directory.
         Args:
             path (str): The directory where the dataset will be downloaded.
-            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):
@@ -74,7 +73,6 @@ 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/metrics/EntropyPred.py

@@ -5,25 +5,32 @@
 class EntropyPrediction(nn.Module):
     def __init__(self, averages: str = "average"):
         """
-        Initializes the EntropyPrediction module.
+        Initializes the EntropyPrediction module, which calculates the Shannon Entropy
+        of predicted logits and aggregates the results based on the specified method.
         Args:
             averages (str): Specifies the method of aggregation for entropy values.
-                            Must be either 'mean', 'sum' or 'none.
+                            Must be one of 'mean', 'sum', or 'none'.
         Raises:
-            AssertionError: If the averages parameter is not 'mean' or 'sum'.
+            AssertionError: If the averages parameter is not 'mean', 'sum', or 'none'.
         """
         super().__init__()
 
-        assert averages == "mean" or averages == "sum"
+        assert averages in ["mean", "sum", "none"], (
+            "averages must be 'mean', 'sum', or 'none'"
+        )
         self.averages = averages
         self.stored_entropy_values = []
 
     def __call__(self, y_true, y_logits):
         """
-        Computes the Shannon Entropy of the predicted logits, storing the results.
+        Computes the Shannon Entropy of the predicted logits and stores the results.
         Args:
-            y_true: The true labels. Does nothing, but needed for compatability sake.
-            y_logits: The predicted logits.
+            y_true: The true labels. This parameter is not used in the computation
+                    but is included for compatibility with certain interfaces.
+            y_logits: The predicted logits from which entropy is calculated.
+        Returns:
+            torch.Tensor: The aggregated entropy value(s) based on the specified
+                          method ('mean', 'sum', or 'none').
         """
         entropy_values = entropy(y_logits, axis=1)
         entropy_values = th.from_numpy(entropy_values)
@@ -34,10 +41,8 @@ def __call__(self, y_true, y_logits):
 
         if self.averages == "mean":
             entropy_values = th.mean(entropy_values)
-
         elif self.averages == "sum":
             entropy_values = th.sum(entropy_values)
-
         elif self.averages == "none":
             return entropy_values
 

utils/models/magnus_model.py

@@ -5,20 +5,16 @@ class MagnusModel(nn.Module):
     def __init__(self, image_shape, num_classes: int, nr_channels: int):
         """
         Initializes the MagnusModel, a neural network designed for image classification tasks.
-
         The model consists of three linear layers, each with 133 neurons, and uses ReLU activation
         functions between the layers. The first layer's input size is determined by the image shape
         and number of channels, while the output layer's size is determined by the number of classes.
         Args:
-            image_shape (tuple): A tuple representing the dimensions of the input image (Channels, Height, Width).
+            image_shape (tuple): A tuple representing the dimensions of the input image (Height, Width).
             num_classes (int): The number of output classes for classification.
             nr_channels (int): The number of channels in the input image.
-        Returns:
-            MagnusModel (nn.Module): An instance of the MagnusModel neural network.
         """
         super().__init__()
         *_, H, W = image_shape
-
         self.layer1 = nn.Sequential(
             *(
                 [
@@ -40,7 +36,8 @@ def forward(self, x):
         """
         Defines the forward pass of the MagnusModel.
         Args:
-            x (torch.Tensor): A four-dimensional tensor with shape (Batch Size, Channels, Image Height, Image Width).
+            x (torch.Tensor): A four-dimensional tensor with shape
+                              (Batch Size, Channels, Image Height, Image Width).
         Returns:
             torch.Tensor: The output tensor containing class logits for each input sample.
         """