Upload model.py

model.py

@@ -0,0 +1,344 @@
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+import math
+
+# class Net(nn.Module):
+#     def __init__(self, upscale_factor):
+#         super(Net, self).__init__()
+#
+#         self.relu = nn.ReLU()
+#         self.conv1 = nn.Conv2d(3, 64, (5, 5), (1, 1), (2, 2))
+#         self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
+#         self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
+#         self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
+#         self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
+#
+#         self._initialize_weights()
+#
+#     def forward(self, x):
+#         x = self.relu(self.conv1(x))
+#         x = self.relu(self.conv2(x))
+#         x = self.relu(self.conv3(x))
+#         x = self.pixel_shuffle(self.conv4(x))
+#         return x
+#
+#     def _initialize_weights(self):
+#         init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
+#         init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
+#         init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
+#         init.orthogonal_(self.conv4.weight)
+
+# class Net(nn.Module):
+#     def __init__(self, upscale_factor=1):
+#         super(Net, self).__init__()
+#
+#         self.relu = nn.ReLU()
+#         self.conv1 = nn.Conv2d(3, 64, (5, 5), (1, 1), (2, 2))
+#         self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
+#         self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
+#         self.conv4 = nn.Conv2d(32, 3, (3, 3), (1, 1), (1, 1))
+#         #self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
+#         self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
+#
+#         self._initialize_weights()
+#
+#     def forward(self, x):
+#         x = self.relu(self.conv1(x))
+#         x = self.relu(self.conv2(x))
+#         x = self.upsample(self.relu(self.conv3(x)))
+#         x = self.relu(self.conv4(x))
+#         return x
+#
+#     def _initialize_weights(self):
+#         init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
+#         init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
+#         init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
+#         init.orthogonal_(self.conv4.weight)
+
+# class Net(nn.Module):
+#     def __init__(self, upscale_factor=1):
+#         super(Net, self).__init__()
+#
+#         self.relu = nn.ReLU()
+#         self.conv1 = nn.Conv2d(3, 64, (5, 5), (1, 1), (2, 2))
+#         self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
+#         self.conv3 = nn.Conv2d(64, 128, (3, 3), (1, 1), (1, 1))
+#         self.conv4 = nn.Conv2d(128, 64, (3, 3), (1, 1), (1, 1))
+#         self.conv5 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
+#         self.conv6 = nn.Conv2d(32, 3, (3, 3), (1, 1), (1, 1))
+#         #self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
+#         self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
+#
+#         self._initialize_weights()
+#
+#
+#     def forward(self, x):
+#         x = self.relu(self.conv1(x))
+#         x = self.relu(self.conv2(x))
+#         x = self.upsample(self.relu(self.conv3(x)))
+#         x = self.relu(self.conv4(x))
+#         x = self.upsample(self.relu(self.conv5(x)))
+#         x = self.relu(self.conv6(x))
+#         return x
+#
+#     def _initialize_weights(self):
+#         init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
+#         init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
+#         init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
+#         init.orthogonal_(self.conv4.weight, init.calculate_gain('relu'))
+#         init.orthogonal_(self.conv5.weight, init.calculate_gain('relu'))
+#         init.orthogonal_(self.conv6.weight, init.calculate_gain('relu'))
+
+class block(nn.Module):
+    def __init__(self,channels_in,channels_out,kernel,stride,pad):
+        super(block, self).__init__()
+        self.conv = nn.Conv2d(channels_in, channels_out, kernel, stride, pad)
+        self.act = nn.LeakyReLU(0.1)
+        self.norm = nn.InstanceNorm2d(channels_out)
+
+    def forward(self, x):
+        x = self.conv(x)
+        #x = self.norm(x)
+        x = self.act(x)
+
+        return x
+
+class SubPixelConvolutionalBlock(nn.Module):
+    """
+    A subpixel convolutional block, comprising convolutional, pixel-shuffle, and PReLU activation layers.
+    """
+
+    def __init__(self, kernel_size=3, n_channels=64, scaling_factor=2):
+        """
+        :param kernel_size: kernel size of the convolution
+        :param n_channels: number of input and output channels
+        :param scaling_factor: factor to scale input images by (along both dimensions)
+        """
+        super(SubPixelConvolutionalBlock, self).__init__()
+
+        # A convolutional layer that increases the number of channels by scaling factor^2, followed by pixel shuffle and PReLU
+        self.conv = nn.Conv2d(in_channels=n_channels, out_channels=n_channels * (scaling_factor ** 2),
+                              kernel_size=kernel_size, padding=kernel_size // 2)
+        # These additional channels are shuffled to form additional pixels, upscaling each dimension by the scaling factor
+        self.pixel_shuffle = nn.PixelShuffle(upscale_factor=scaling_factor)
+        self.lrelu = nn.LeakyReLU(0.1)
+
+    def forward(self, input):
+        """
+        Forward propagation.
+
+        :param input: input images, a tensor of size (N, n_channels, w, h)
+        :return: scaled output images, a tensor of size (N, n_channels, w * scaling factor, h * scaling factor)
+        """
+        output = self.conv(input)  # (N, n_channels * scaling factor^2, w, h)
+        output = self.pixel_shuffle(output)  # (N, n_channels, w * scaling factor, h * scaling factor)
+        output = self.lrelu(output)  # (N, n_channels, w * scaling factor, h * scaling factor)
+
+        return output
+
+class Net(nn.Module):
+    def __init__(self, upscale_factor=1):
+        super(Net, self).__init__()
+
+        self.conv1 = block(3, 64, (7, 7), (1, 1), (3, 3))
+        self.conv2 = block(64, 64, (5, 5), (1, 1), (2, 2))
+        self.conv3 = block(64, 128, (5, 5), (1, 1), (2, 2))
+        self.conv4 = block(128, 256, (5, 5), (1, 1), (2, 2))
+        self.conv4_1 = block(256, 256, (5, 5), (1, 1), (2, 2))
+        self.conv5 = block(256, 128, (5, 5), (1, 1), (2, 2))
+        self.conv6 = block(128, 64, (3, 3), (1, 1), (1, 1))
+        self.conv7 = block(64, 32, (3, 3), (1, 1), (1, 1))
+        self.conv7_1 = block(32, 16, (3, 3), (1, 1), (1, 1))
+        self.conv8 = nn.Conv2d(16, 3, (1, 1), (1, 1), (0, 0))
+        #self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
+        self.relu = nn.ReLU()
+        #self.upsample = nn.Upsample(scale_factor=2, mode='bicubic')
+        self.spc1 = SubPixelConvolutionalBlock( kernel_size=3, n_channels=128, scaling_factor=2)
+        #self.spc1 = nn.Upsample(scale_factor=2, mode='bicubic')
+        self.spc2 = SubPixelConvolutionalBlock( kernel_size=3, n_channels=128, scaling_factor=2)
+        #self.spc2 = nn.Upsample(scale_factor=2, mode='bicubic')
+
+        self.spc3 = SubPixelConvolutionalBlock(kernel_size=3, n_channels=128, scaling_factor=4)
+
+        #self._initialize_weights()
+
+    def forward(self, x):
+        x =self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.spc1(x)
+        x = self.conv4(x)
+        x = self.conv4_1(x)
+        x = self.conv5(x)
+        x = self.spc2(x)
+        x = self.conv6(x)
+        x = self.conv7(x)
+        x = self.conv7_1(x)
+        x = self.relu(self.conv8(x))
+        return x
+
+    def _initialize_weights(self):
+        init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
+        init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
+        init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
+        init.orthogonal_(self.conv4.weight, init.calculate_gain('relu'))
+        init.orthogonal_(self.conv4_1.weight, init.calculate_gain('relu'))
+        init.orthogonal_(self.conv5.weight, init.calculate_gain('relu'))
+        init.orthogonal_(self.conv6.weight, init.calculate_gain('relu'))
+        init.orthogonal_(self.conv7.weight, init.calculate_gain('relu'))
+        init.orthogonal_(self.conv7_1.weight, init.calculate_gain('relu'))
+        init.orthogonal_(self.conv8.weight, init.calculate_gain('relu'))
+
+
+
+
+
+class ConvolutionalBlock(nn.Module):
+    """
+    A convolutional block, comprising convolutional, BN, activation layers.
+    """
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, batch_norm=False, activation=None):
+        """
+        :param in_channels: number of input channels
+        :param out_channels: number of output channe;s
+        :param kernel_size: kernel size
+        :param stride: stride
+        :param batch_norm: include a BN layer?
+        :param activation: Type of activation; None if none
+        """
+        super(ConvolutionalBlock, self).__init__()
+
+        if activation is not None:
+            activation = activation.lower()
+            assert activation in {'prelu', 'leakyrelu', 'tanh'}
+
+        # A container that will hold the layers in this convolutional block
+        layers = list()
+
+        # A convolutional layer
+        layers.append(
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride,
+                      padding=kernel_size // 2))
+
+        # A batch normalization (BN) layer, if wanted
+        if batch_norm is True:
+            layers.append(nn.BatchNorm2d(num_features=out_channels))
+
+        # An activation layer, if wanted
+        if activation == 'prelu':
+            layers.append(nn.PReLU())
+        elif activation == 'leakyrelu':
+            layers.append(nn.LeakyReLU(0.2))
+        elif activation == 'tanh':
+            layers.append(nn.Tanh())
+
+        # Put together the convolutional block as a sequence of the layers in this container
+        self.conv_block = nn.Sequential(*layers)
+
+    def forward(self, input):
+        """
+        Forward propagation.
+
+        :param input: input images, a tensor of size (N, in_channels, w, h)
+        :return: output images, a tensor of size (N, out_channels, w, h)
+        """
+        output = self.conv_block(input)  # (N, out_channels, w, h)
+
+        return output
+
+
+class ResidualBlock(nn.Module):
+    """
+    A residual block, comprising two convolutional blocks with a residual connection across them.
+    """
+
+    def __init__(self, kernel_size=3, n_channels=64):
+        """
+        :param kernel_size: kernel size
+        :param n_channels: number of input and output channels (same because the input must be added to the output)
+        """
+        super(ResidualBlock, self).__init__()
+
+        # The first convolutional block
+        self.conv_block1 = ConvolutionalBlock(in_channels=n_channels, out_channels=n_channels, kernel_size=kernel_size,
+                                              batch_norm=True, activation='PReLu')
+
+        # The second convolutional block
+        self.conv_block2 = ConvolutionalBlock(in_channels=n_channels, out_channels=n_channels, kernel_size=kernel_size,
+                                              batch_norm=True, activation=None)
+
+    def forward(self, input):
+        """
+        Forward propagation.
+
+        :param input: input images, a tensor of size (N, n_channels, w, h)
+        :return: output images, a tensor of size (N, n_channels, w, h)
+        """
+        residual = input  # (N, n_channels, w, h)
+        output = self.conv_block1(input)  # (N, n_channels, w, h)
+        output = self.conv_block2(output)  # (N, n_channels, w, h)
+        output = output + residual  # (N, n_channels, w, h)
+
+        return output
+
+
+class SRResNet(nn.Module):
+    """
+    The SRResNet, as defined in the paper.
+    """
+
+    def __init__(self, large_kernel_size=9, small_kernel_size=3, n_channels=64, n_blocks=16, scaling_factor=4):
+        """
+        :param large_kernel_size: kernel size of the first and last convolutions which transform the inputs and outputs
+        :param small_kernel_size: kernel size of all convolutions in-between, i.e. those in the residual and subpixel convolutional blocks
+        :param n_channels: number of channels in-between, i.e. the input and output channels for the residual and subpixel convolutional blocks
+        :param n_blocks: number of residual blocks
+        :param scaling_factor: factor to scale input images by (along both dimensions) in the subpixel convolutional block
+        """
+        super(SRResNet, self).__init__()
+
+        # Scaling factor must be 2, 4, or 8
+        scaling_factor = int(scaling_factor)
+        assert scaling_factor in {2, 4, 8}, "The scaling factor must be 2, 4, or 8!"
+
+        # The first convolutional block
+        self.conv_block1 = ConvolutionalBlock(in_channels=3, out_channels=n_channels, kernel_size=large_kernel_size,
+                                              batch_norm=False, activation='PReLu')
+
+        # A sequence of n_blocks residual blocks, each containing a skip-connection across the block
+        self.residual_blocks = nn.Sequential(
+            *[ResidualBlock(kernel_size=small_kernel_size, n_channels=n_channels) for i in range(n_blocks)])
+
+        # Another convolutional block
+        self.conv_block2 = ConvolutionalBlock(in_channels=n_channels, out_channels=n_channels,
+                                              kernel_size=small_kernel_size,
+                                              batch_norm=True, activation=None)
+
+        # Upscaling is done by sub-pixel convolution, with each such block upscaling by a factor of 2
+        n_subpixel_convolution_blocks = int(math.log2(scaling_factor))
+        self.subpixel_convolutional_blocks = nn.Sequential(
+            *[SubPixelConvolutionalBlock(kernel_size=small_kernel_size, n_channels=n_channels, scaling_factor=2) for i
+              in range(n_subpixel_convolution_blocks)])
+
+        # The last convolutional block
+        self.conv_block3 = ConvolutionalBlock(in_channels=n_channels, out_channels=3, kernel_size=large_kernel_size,
+                                              batch_norm=False, activation='Tanh')
+
+    def forward(self, lr_imgs):
+        """
+        Forward prop.
+
+        :param lr_imgs: low-resolution input images, a tensor of size (N, 3, w, h)
+        :return: super-resolution output images, a tensor of size (N, 3, w * scaling factor, h * scaling factor)
+        """
+        output = self.conv_block1(lr_imgs)  # (N, 3, w, h)
+        residual = output  # (N, n_channels, w, h)
+        output = self.residual_blocks(output)  # (N, n_channels, w, h)
+        output = self.conv_block2(output)  # (N, n_channels, w, h)
+        output = output + residual  # (N, n_channels, w, h)
+        output = self.subpixel_convolutional_blocks(output)  # (N, n_channels, w * scaling factor, h * scaling factor)
+        sr_imgs = self.conv_block3(output)  # (N, 3, w * scaling factor, h * scaling factor)
+
+        return sr_imgs