Add the changes to comply with the project specifications. (#31)

* Added a first implementation of conv2d

* Added the non-optimized implementation of conv2d

We maintained the same names of attributes to make it compatible with the checkpoints.

* Implemented a simulation of random_split of pytorch

* Implemented a simulation of PixelShuffle module

Given the previous layer return the right dimensions of the tensor, the pixel shuffle is just a reshaping.

* Changed the bias to a tensor of zeros when bias=False

This change is just for get rid of the warning that int has no __get_item__ implementation

* Added a im2col fast convolution and fixed PixelShuffle

* Added small model training

It only upscales 32x16 ==> 64x32

* Changed constructor of SuperResolutionDataset

Added 2 parameters to change the low and high resolution of the loading. Default values are 128x64 ==> 256x128

* Added declaration for small dataset training

* Added old validation parameters

* Added in the notebook the comparison between models

* Changed notebook date generation

* Added notebook plots

* Added generalization to choose resolutions

* Changed return statement

* Edited the report to include both models

* Deleted old files

README.md

@@ -38,7 +38,7 @@ python3 main.py
 # Export the notebook as pdf
 
 ```bash
-jupyter nbconvert --to pdf notebook.ipynb --output "ModelDemonstration" --LatexPreprocessor.title "Super Resolution Demonstration" --LatexPreprocessor.date "August, 2024" --LatexPreprocessor.author_names "Christian Mancini"
+jupyter nbconvert --to pdf notebook.ipynb --output "ModelDemonstration" --LatexPreprocessor.title "Super Resolution Demonstration" --LatexPreprocessor.date "September, 2024" --LatexPreprocessor.author_names "Christian Mancini"
 ```
 
 # Installation of Requirements and Kernel

SRM/modules.py

@@ -1,13 +1,30 @@
+import torch
 from torch import nn
+import numpy as np
 
 
+def _im2col(input, kernel_size, stride=1, padding=0):
+    input_padded = torch.nn.functional.pad(input, (padding, padding, padding, padding))
+    batch_size, in_channels, height, width = input_padded.size()
+    kernel_height, kernel_width = kernel_size
+    out_height = (height - kernel_height) // stride + 1
+    out_width = (width - kernel_width) // stride + 1
+    col = torch.empty(batch_size, in_channels, kernel_height, kernel_width, out_height, out_width)
+
+    for y in range(kernel_height):
+        for x in range(kernel_width):
+            col[:, :, y, x, :, :] = input_padded[:, :, y: y + out_height * stride: stride,
+                                    x: x + out_width * stride: stride]
+
+    return col.view(batch_size, in_channels * kernel_height * kernel_width, -1)
+
 class ResidualBlock(nn.Module):
 
     def __init__(self, num_channels):
         super().__init__()
-        self.conv1 = nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1)
+        self.conv1 = Conv2d(num_channels, num_channels, kernel_size=3, padding=1)
         self.relu = nn.ReLU()
-        self.conv2 = nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1)
+        self.conv2 = Conv2d(num_channels, num_channels, kernel_size=3, padding=1)
 
     def forward(self, x):
         res = x
@@ -22,11 +39,80 @@ class Upsample(nn.Module):
 
     def __init__(self, num_channel):
         super().__init__()
-        self.conv = nn.Conv2d(num_channel, num_channel * 4, kernel_size=3, padding=1)
-        self.shuffle = nn.PixelShuffle(2)
+        self.conv = Conv2d(num_channel, num_channel * 4, kernel_size=3, padding=1)
+        self.shuffle = PixelShuffle(2)
 
     def forward(self, x):
         out = self.conv(x)
         out = self.shuffle(out)
         return out
 
+
+class Conv2d(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, padding=1, bias=True):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.padding = padding
+        self.k = 1/(in_channels * np.power(kernel_size, 2))
+
+        if bias:
+            self.bias = nn.Parameter(
+                torch.empty(out_channels).uniform_(-np.sqrt(self.k),np.sqrt(self.k)))
+        else:
+            self.bias = torch.zeros(out_channels)
+
+        self.weight = nn.Parameter(
+            torch.empty(
+                out_channels, in_channels, kernel_size, kernel_size).uniform_(-np.sqrt(self.k),np.sqrt(self.k)))
+
+
+    def _conv_forward(self, input, weight, bias=None, stride=1, padding=0):
+        col = _im2col(input, weight.size()[2:], stride, padding)
+        # (out_channels, in_channels * kernel_height * kernel_width)
+        weight_col = weight.view(weight.size(0), -1)
+        out = torch.matmul(weight_col, col)
+
+        if bias is not None:
+            out += bias.view(1, -1, 1)
+
+        batch_size, out_channels = out.size(0), weight.size(0)
+        out_height = (input.size(2) + 2 * padding - weight.size(2)) // stride + 1
+        out_width = (input.size(3) + 2 * padding - weight.size(3)) // stride + 1
+        return out.view(batch_size, out_channels, out_height, out_width)
+
+    def forward(self, input):
+        return self._conv_forward(input, self.weight, self.bias, padding=self.padding)
+
+
+    def slow_forward(self, image):
+        image = nn.functional.pad(image, (self.padding,) * 4, "constant", 0)
+        batch_size, in_channels, height, width = image.shape
+        out_channels, in_channels_kernel, m, n = self.weight.shape
+        if self.in_channels != in_channels:
+            raise ValueError(
+                f"Input channels are different: Declared {self.in_channels}, but got Image with {in_channels}")
+        output_height = height - m + 1
+        output_width = width - n + 1
+        new_image = torch.zeros((batch_size, out_channels, output_height, output_width))
+
+        for b in range(batch_size):
+            for c in range(out_channels):
+                for i in range(output_height):
+                    for j in range(output_width):
+                        new_image[b, c, i, j] = torch.sum(image[b, :, i:i + m, j:j + n] * self.weight[c]) + self.bias[c]
+        return new_image
+
+
+class PixelShuffle(nn.Module):
+    def __init__(self, upscale_factor: int):
+        super().__init__()
+        self.upscale_factor = upscale_factor
+
+    def forward(self, x):
+        batch_size, channels, height, width = x.shape
+        channels //= (self.upscale_factor ** 2)
+        x = x.view(batch_size, channels, self.upscale_factor, self.upscale_factor, height, width)
+        x = x.permute(0, 1, 4, 2, 5, 3)
+        return x.contiguous().view(batch_size, channels, height * self.upscale_factor, width * self.upscale_factor)

SRM/network.py

@@ -9,12 +9,12 @@ class SuperResolution(nn.Module):
     def __init__(self, num_channels: int, num_res_block: int):
         super().__init__()
         self.layers = nn.ModuleList([
-            nn.Conv2d(3, num_channels, kernel_size=3, padding=1)
+            Conv2d(3, num_channels, kernel_size=3, padding=1)
         ])
         for _ in range(num_res_block):
             self.layers.append(ResidualBlock(num_channels))
         self.layers.append(Upsample(num_channels))
-        self.layers.append(nn.Conv2d(num_channels, 3, kernel_size=3, padding=1))
+        self.layers.append(Conv2d(num_channels, 3, kernel_size=3, padding=1))
 
     def forward(self, x):
         for layer in self.layers:

dataset/data_preparation.py

@@ -29,6 +29,28 @@ def download(path: str, dataset: str):
         print(f"Dataset {dataset} already exists, skipping download.")
 
 
+def _random_split(dataset: Dataset, lengths: list[int]) -> list[Subset]:
+    """
+    Utility method that simulate pytorch random split
+    Args:
+        dataset: Dataset to be split
+        lengths: List of lengths to split the dataset into
+
+    Returns:
+        List of Subsets
+    """
+    if sum(lengths) != len(dataset):
+        raise ValueError("Sum of input lengths does not equal the length of the input dataset!")
+
+    indices = torch.randperm(len(dataset)).tolist()
+    subsets = []
+    offset = 0
+    for length in lengths:
+        subset = torch.utils.data.Subset(dataset, indices[offset:offset + length])
+        subsets.append(subset)
+        offset += length
+    return subsets
+
 def split_dataset(dataset: Dataset, sizes: dict[str, float]) -> list[Subset]:
     """
         Splits a dataset into training, validation, and test sets based on the provided sizes.
@@ -52,4 +74,4 @@ def split_dataset(dataset: Dataset, sizes: dict[str, float]) -> list[Subset]:
     train_size = int(sizes["train"] * len(dataset))
     validation_size = int(sizes["validation"] * len(dataset))
     test_size = len(dataset) - train_size - validation_size
-    return torch.utils.data.random_split(dataset, [train_size, validation_size, test_size])
+    return _random_split(dataset, [train_size, validation_size, test_size])

dataset/super_resolution_dataset.py

@@ -6,10 +6,13 @@
 
 
 class SuperResolutionDataset(Dataset):
-    def __init__(self, root_dir, transform=transforms.Compose([transforms.ToTensor()])) -> None:
+    def __init__(self, root_dir, low_resolution=(128, 64), high_resolution=(256, 128),
+                 transform=transforms.Compose([transforms.ToTensor()])) -> None:
         self.root_dir = root_dir
         self.transform = transform
         self.file_names = os.listdir(root_dir)
+        self.low_resolution = low_resolution
+        self.high_resolution = high_resolution
 
     def __len__(self) -> int:
         return len(self.file_names)
@@ -25,8 +28,8 @@ def __getitem__(self, idx) -> tuple[torch.Tensor, torch.Tensor]:
         """
         img_path = os.path.join(self.root_dir, self.file_names[idx])
         image = Image.open(img_path)
-        low_res = image.resize((128, 64))
-        high_res = image.resize((256, 128))
+        low_res = image.resize(self.low_resolution)
+        high_res = image.resize(self.high_resolution)
 
         if self.transform:
             low_res = self.transform(low_res)

main.py

@@ -23,6 +23,9 @@
 training_epochs = 50
 final_training_epochs = 150
 
+# Define low and high resolution for loading data and training
+low_res = (32,16)
+high_res = (64,32)
 
 def main():
     torch.manual_seed(777)
@@ -33,7 +36,7 @@ def main():
     # Download and extract the dataset
     download("./data", "airplanes")
     root_dir = 'data/airplanes'
-    dataset = SuperResolutionDataset(root_dir=root_dir)
+    dataset = SuperResolutionDataset(root_dir=root_dir,low_resolution=low_res,high_resolution=high_res)
 
     # Split the dataset and make the final training dataset (to be used after model selection)
     train, validation, test = split_dataset(dataset, sizes)