Added python docs

DanielHalachev · DanielHalachev · commit 2645aef9a448 · 2024-01-28T10:09:18.000+02:00
diff --git a/dataset.py b/dataset.py
@@ -8,7 +8,18 @@
 
 
 class ColorizationDataset(Dataset):
+    """
+    Custom class for loading the dataset into the neural network using PyTorch
+    """
+
     def __init__(self, paths, split='train'):
+        """
+        Colorization Dataset initializer
+        :param paths: List of file paths for images
+        :type paths: list
+        :param split:Dataset split, either 'train' or 'test'
+        :type split: str, optional
+        """
         if split == 'train':
             self.transforms = transforms.Compose([
                 transforms.Resize((SIZE, SIZE), Image.BICUBIC),
@@ -37,6 +48,20 @@ def __len__(self):
 
 
 def make_dataloaders(batch_size=16, n_workers=4, pin_memory=True, **kwargs):
+    """
+        Create DataLoader for ColorizationDataset.
+
+        :param batch_size: Number of samples per batch. Default is 16.
+        :type batch_size: int, optional
+        :param n_workers: Number of workers for data loading. Default is 4.
+        :type n_workers: int, optional
+        :param pin_memory: Whether to use pinned memory for faster data transfer. Default is True.
+        :type pin_memory: bool, optional
+        :param kwargs: Additional arguments to pass to ColorizationDataset.
+
+        :return: DataLoader for the ColorizationDataset.
+        :rtype: torch.utils.data.DataLoader
+    """
     dataset = ColorizationDataset(**kwargs)
     dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=n_workers, pin_memory=pin_memory)
     return dataloader
diff --git a/loss.py b/loss.py
@@ -3,7 +3,18 @@
 
 
 class GANLoss(nn.Module):
+    """
+    A module for measuring the GAN Loss.
+    """
     def __init__(self, gan_mode='vanilla', real_label=1.0, fake_label=0.0):
+        """
+        Initializes the GANLoss module.
+
+        :param real_label: The label value for real samples (default is 1.0).
+        :type real_label: float
+        :param fake_label: The label value for fake/generated samples (default is 0.0).
+        :type fake_label: float
+        """
         super().__init__()
         self.register_buffer('real_label', torch.tensor(real_label))
         self.register_buffer('fake_label', torch.tensor(fake_label))
@@ -13,13 +24,35 @@ def __init__(self, gan_mode='vanilla', real_label=1.0, fake_label=0.0):
             self.loss = nn.MSELoss()
 
     def get_labels(self, preds, target_is_real):
+        """
+        Returns the target labels based on whether the target is real or fake.
+
+        :param preds:The predictions from the discriminator.
+        :type preds: torch.Tensor
+        :param target_is_real:Indicates whether the target is a real sample.
+        :type target_is_real: bool
+
+        :return: The target labels expanded to match the shape of predictions.
+        :rtype: torch.Tensor
+        """
         if target_is_real:
             labels = self.real_label
         else:
             labels = self.fake_label
         return labels.expand_as(preds)
 
     def __call__(self, preds, target_is_real):
+        """
+        Computes and returns the adversarial loss given the predictions and target labels.
+
+        :param preds: The predictions from the discriminator.
+        :type preds: torch.Tensor
+        :param target_is_real: Indicates whether the target is a real sample.
+        :type target_is_real: bool
+
+        :return: The computed adversarial loss.
+        :rtype: torch.Tensor
+        """
         labels = self.get_labels(preds, target_is_real)
         loss = self.loss(preds, labels)
         return loss
diff --git a/models.py b/models.py
@@ -4,7 +4,28 @@
 
 
 class UnetBlock(nn.Module):
+    """
+        Defines a U-Net block, a fundamental building block for the generator in an image-to-image translation model.
+    """
     def __init__(self, nf, ni, submodule=None, input_channels=None, dropout=False, innermost=False, outermost=False):
+        """
+        Initialize UnetBlock.
+
+        :param nf: Number of filters.
+        :type nf: int
+        :param ni: Number of input channels.
+        :type ni: int
+        :param submodule: Submodule to be included inside the block. Default is None.
+        :type submodule: nn.Module, optional
+        :param input_channels: Number of input channels. Default is None.
+        :type input_channels: int, optional
+        :param dropout: Whether to apply dropout. Default is False.
+        :type dropout: bool, optional
+        :param innermost: Whether the block is innermost. Default is False.
+        :type innermost: bool, optional
+        :param outermost: Whether the block is outermost. Default is False.
+        :type outermost: bool, optional
+        """
         super().__init__()
         self.outermost = outermost
         if input_channels is None:
@@ -35,13 +56,25 @@ def __init__(self, nf, ni, submodule=None, input_channels=None, dropout=False, i
         self.model = nn.Sequential(*model)
 
     def forward(self, x):
+        """
+        Forward pass through the U-Net block.
+
+        :param x: Input tensor.
+        :type x: torch.Tensor
+
+        :return: Output tensor.
+        :rtype: torch.Tensor
+        """
         if self.outermost:
             return self.model(x)
         else:
             return torch.cat([x, self.model(x)], 1)
 
 
 class Unet(nn.Module):
+    """
+    Defines a U-Net model constructed using U-Net Blocks
+    """
     def __init__(self, nfg=64):
         super().__init__()
         unet_block = UnetBlock(nfg * 8, nfg * 8, innermost=True)
@@ -54,11 +87,29 @@ def __init__(self, nfg=64):
         self.model = UnetBlock(2, out_filters, input_channels=1, submodule=unet_block, outermost=True)
 
     def forward(self, x):
+        """
+        Forward pass through the U-Net model.
+
+        :param x: Input tensor.
+        :type x: torch.Tensor
+
+        :return: Output tensor.
+        :rtype: torch.Tensor
+        """
         return self.model(x)
 
 
 class PatchDiscriminator(nn.Module):
+    """
+    Defines a PatchGAN discriminator for image-to-image translation tasks.
+    """
     def __init__(self, nfd=64):
+        """
+        Initialize PatchDiscriminator.
+
+        :param nfd: Number of initial filters. Default is 64.
+        :type nfd: int, optional
+        """
         super().__init__()
         # No normalization in first block
         model = [self.get_layers(3, nfd, normalization=False)]
@@ -70,6 +121,27 @@ def __init__(self, nfd=64):
         self.model = nn.Sequential(*model)
 
     def get_layers(self, ni, nf, k=4, s=2, p=1, normalization=True, action=True):
+        """
+        Helper method to create layers for the PatchGAN discriminator.
+
+        :param ni: Number of input channels.
+        :type ni: int
+        :param nf: Number of filters.
+        :type nf: int
+        :param k: Kernel size. Default is 4.
+        :type k: int, optional
+        :param s: Stride. Default is 2.
+        :type s: int, optional
+        :param p: Padding. Default is 1.
+        :type p: int, optional
+        :param normalization: Whether to apply batch normalization. Default is True.
+        :type normalization: bool, optional
+        :param action: Whether to apply activation function. Default is True.
+        :type action: bool, optional
+
+        :return: Sequential model representing the layers.
+        :rtype: nn.Sequential
+        """
         layers = [
             nn.Conv2d(ni, nf, k, s, p, bias=not normalization)]
         if normalization:
@@ -79,10 +151,32 @@ def get_layers(self, ni, nf, k=4, s=2, p=1, normalization=True, action=True):
         return nn.Sequential(*layers)
 
     def forward(self, x):
+        """
+        Forward pass through the PatchGAN discriminator.
+
+        :param x: Input tensor.
+        :type x: torch.Tensor
+
+        :return: Output tensor.
+        :rtype: torch.Tensor
+        """
         return self.model(x)
 
 
 def init_weights(net, init='norm', gain=0.02):
+    """
+    Initialize weights for the neural network.
+
+    :param net: Neural network model.
+    :type net: nn.Module
+    :param init: Initialization method. Default is 'norm'.
+    :type init: str, optional
+    :param gain: Gain factor for weight initialization. Default is 0.02.
+    :type gain: float, optional
+
+    :return: Initialized neural network model.
+    :rtype: nn.Module
+    """
     def init_func(m):
         classname = m.__class__.__name__
         if hasattr(m, 'weight') and 'Conv' in classname:
@@ -105,19 +199,49 @@ def init_func(m):
 
 
 def init_model(model, device):
+    """
+    Initialize a neural network model.
+
+    :param model: Neural network model.
+    :type model: nn.Module
+    :param device: Device to which the model will be moved.
+    :type device: torch.device
+
+    :return: Initialized and moved neural network model.
+    :rtype: nn.Module
+    """
     model = model.to(device)
     model = init_weights(model)
     return model
 
 
 class MainModel(nn.Module):
+    """
+    Main model for image-to-image translation tasks using a conditional GAN with a U-Net generator.
+    """
     def __init__(self,
                  net_generator=None,
                  lr_generator=2e-4,
                  lr_discriminator=2e-4,
                  beta1=0.5,
                  beta2=0.999,
                  lambda_l1=100.):
+        """
+        Initialize MainModel.
+
+        :param net_generator: Predefined generator network. Default is None.
+        :type net_generator: nn.Module, optional
+        :param lr_generator: Learning rate for the generator. Default is 2e-4.
+        :type lr_generator: float, optional
+        :param lr_discriminator: Learning rate for the discriminator. Default is 2e-4.
+        :type lr_discriminator: float, optional
+        :param beta1: Beta1 parameter for Adam optimizer. Default is 0.5.
+        :type beta1: float, optional
+        :param beta2: Beta2 parameter for Adam optimizer. Default is 0.999.
+        :type beta2: float, optional
+        :param lambda_l1: Weight for L1 loss term. Default is 100.
+        :type lambda_l1: float, optional
+        """
         super().__init__()
 
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -134,17 +258,37 @@ def __init__(self,
         self.opt_D = optim.Adam(self.net_D.parameters(), lr=lr_discriminator, betas=(beta1, beta2))
 
     def set_requires_grad(self, model, requires_grad=True):
+        """
+        Set the requires_grad attribute for model parameters.
+
+        :param model: Model for which to set the requires_grad attribute.
+        :type model: nn.Module
+        :param requires_grad: Whether to set requires_grad to True or False. Default is True.
+        :type requires_grad: bool, optional
+        """
         for p in model.parameters():
             p.requires_grad = requires_grad
 
     def setup_input(self, data):
+        """
+        Move input data to the specified device.
+
+        :param data: Input data containing LAB colorspace components.
+        :type data: dict
+        """
         self.L = data['L'].to(self.device)
         self.ab = data['ab'].to(self.device)
 
     def forward(self):
+        """
+        Forward pass through the generator
+        """
         self.fake_color = self.net_G(self.L)
 
     def backward_D(self):
+        """
+        Backward pass and optimization for the discriminator.
+        """
         fake_image = torch.cat([self.L, self.fake_color], dim=1)
         fake_predictions = self.net_D(fake_image.detach())
         self.loss_D_fake = self.GAN_criterion(fake_predictions, False)
@@ -155,6 +299,9 @@ def backward_D(self):
         self.loss_D.backward()
 
     def backward_G(self):
+        """
+        Backward pass for the generator.
+        """
         fake_image = torch.cat([self.L, self.fake_color], dim=1)
         fake_predictions = self.net_D(fake_image)
         self.loss_G_GAN = self.GAN_criterion(fake_predictions, True)
@@ -163,6 +310,20 @@ def backward_G(self):
         self.loss_G.backward()
 
     def optimize(self):
+        """
+        Optimization of the whole model.
+
+        This method performs a single optimization step for both the generator and the discriminator.
+        It includes the forward pass, backward pass, and parameter updates.
+
+        Steps:
+        1. Perform the forward pass through the generator.
+        2. Set the discriminator to training mode and enable gradients for its parameters.
+        3. Zero the gradients of the discriminator optimizer.
+        4. Perform the backward pass for the discriminator and update its parameters.
+        5. Set the generator to training mode and disable gradients for the discriminator parameters.
+        6. Zero the gradients of the generator optimizer.
+        """
         self.forward()
         self.net_D.train()
         self.set_requires_grad(self.net_D, True)
diff --git a/train.py b/train.py
@@ -1,7 +1,21 @@
 
 
 def train_model(model, train_dataloader, epochs, display_every=100):
-    # visualization_data = next(iter(val_dl))  # getting a batch for visualizing the model output after fixed intervals
+    """
+    Train the specified model using the provided data loader.
+
+    :param model: (nn.Module) The PyTorch model to be trained.
+    :type model: nn.Module
+    :param train_dataloader: (DataLoader) The DataLoader for training data.
+    :type train_dataloader: DataLoader
+    :param epochs: (int) The number of training epochs.
+    :type epochs: int
+    :param display_every: (int) Display training progress every specified number of iterations (default is 100).
+    :type display_every: int
+
+    Example:
+        train_model(my_model, train_data_loader, epochs=10, display_every=50)
+    """
     for e in range(epochs):
         loss_meter_dict = create_loss_meters()
         i = 0
diff --git a/utils.py b/utils.py
