Upload model.py

model.py

@@ -90,11 +90,76 @@ import math
 #         init.orthogonal_(self.conv5.weight, init.calculate_gain('relu'))
 #         init.orthogonal_(self.conv6.weight, init.calculate_gain('relu'))
 
+class ResidualBlockG(nn.Module):
+  def __init__(self, channels):
+    super(ResidualBlockG, self).__init__()
+    self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+    self.bn1 = nn.BatchNorm2d(channels)
+    self.prelu = nn.ReLU()
+    self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+    self.bn2 = nn.BatchNorm2d(channels)
+  def forward(self, x):
+    residual = self.conv1(x)
+    residual = self.bn1(residual)
+    residual = self.prelu(residual)
+    residual = self.conv2(residual)
+    residual = self.bn2(residual)
+    return x + residual
+
+class UpsampleBlock(nn.Module):
+  def __init__(self, in_channels, up_scale):
+    super(UpsampleBlock, self).__init__()
+    self.conv = nn.Conv2d(in_channels, in_channels * up_scale ** 2,
+                          kernel_size=3, padding=1)
+    self.pixel_shuffle = nn.PixelShuffle(up_scale)
+    self.prelu = nn.ReLU()
+  def forward(self, x):
+    x = self.conv(x)
+    x = self.pixel_shuffle(x)
+    x = self.prelu(x)
+    return x
+
+
+class Generator(nn.Module):
+    def __init__(self, scale_factor):
+        super(Generator, self).__init__()
+        upsample_block_num = int(math.log(scale_factor, 2))
+
+        self.block1 = nn.Sequential(
+            nn.Conv2d(3, 64, kernel_size=9, padding=4),
+            nn.ReLU()
+        )
+
+        self.block2 = ResidualBlockG(64)
+        self.block3 = ResidualBlockG(64)
+        self.block4 = ResidualBlockG(64)
+        self.block5 = ResidualBlockG(64)
+        self.block6 = ResidualBlockG(64)
+        self.block7 = nn.Sequential(
+            nn.Conv2d(64, 64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64)
+        )
+        block8 = [UpsampleBlock(64, 2) for _ in range(upsample_block_num)]
+        block8.append(nn.Conv2d(64, 3, kernel_size=9, padding=4))
+        self.block8 = nn.Sequential(*block8)
+
+    def forward(self, x):
+        block1 = self.block1(x)
+        block2 = self.block2(block1)
+        block3 = self.block3(block2)
+        block4 = self.block4(block3)
+        block5 = self.block5(block4)
+        block6 = self.block6(block5)
+        block7 = self.block7(block6)
+        block8 = self.block8(block1 + block7)
+        return block8
+
 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.act = nn.LeakyReLU(0.02)
+        self.act = nn.ReLU()
         self.norm = nn.InstanceNorm2d(channels_out)
 
     def forward(self, x):
@@ -122,7 +187,7 @@ class SubPixelConvolutionalBlock(nn.Module):
                               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)
+        self.lrelu = nn.ReLU()
 
     def forward(self, input):
         """
@@ -152,13 +217,12 @@ class Net(nn.Module):
         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.relu = nn.ReLU(inplace=True)
         #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()
@@ -178,21 +242,118 @@ class Net(nn.Module):
         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'))
+
     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'))
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                # n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                # m.weight.data.normal_(0, math.sqrt(2. / n))
+                # g = nn.init.calculate_gain('leaky_relu', 0.1)
+                m.weight.data.normal_(0, 0.01)
+                # nn.init.xavier_normal_(m.weight,gai=g)
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.InstanceNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                m.weight.data.normal_(0, 0.01)
+                m.bias.data.zero_()
+        # self.sp.weight.data.normal_(0,0.01)
+
+
+class dws_block(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3, padding=1):
+        super(dws_block, self).__init__()
+        self.dc = nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, padding=padding, groups=in_channels)
+        self.pc = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)
 
+    def forward(self, x):
+        # Encoder
+        x = self.dc(x)
+        x = self.pc(x)
+
+        return x
 
 
+class UNet(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(UNet, self).__init__()
+
+        # Encoder (contracting path)
+        self.encoder1 = self.contracting_block(in_channels, 64)
+        self.encoder2 = self.contracting_block(64, 128)
+        self.encoder3 = self.contracting_block(128, 256)
+        self.encoder4 = self.contracting_block(256, 512)
+
+        # Bottleneck
+        self.bottleneck = nn.Sequential(
+            dws_block(512, 1024, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            dws_block(1024, 1024, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True)
+        )
+
+        # Decoder (expansive path)
+        self.decoder1 = self.expansive_block(1024, 512)
+        self.decoder2 = self.expansive_block(512, 256)
+        self.decoder3 = self.expansive_block(256, 128)
+        self.decoder4 = self.expansive_block(128, 64)
+        self.decoder5 = self.expansive_block(32, 32)
+        self.decoder6 = self.expansive_block(16, 32)
+
+        # Output layer
+        self.final_conv = nn.Conv2d(16, out_channels, kernel_size=1)
+
+    def contracting_block(self, in_channels, out_channels):
+        return nn.Sequential(
+            dws_block(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            dws_block(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+
+    def expansive_block(self, in_channels, out_channels):
+        return nn.Sequential(
+            dws_block(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            dws_block(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(out_channels, out_channels // 2, kernel_size=2, stride=2)
+        )
 
+    def forward(self, x):
+        # Encoder
+        enc1 = self.encoder1(x)
+        enc2 = self.encoder2(enc1)
+        enc3 = self.encoder3(enc2)
+        enc4 = self.encoder4(enc3)
+
+        # Bottleneck
+        bottleneck = self.bottleneck(enc4)
+
+        # Decoder
+        dec1 = self.decoder1(bottleneck)
+        dec2 = self.decoder2(torch.cat([dec1, enc3], dim=1))
+        dec3 = self.decoder3(torch.cat([dec2, enc2], dim=1))
+        dec4 = self.decoder4(torch.cat([dec3, enc1], dim=1))
+        dec5 = self.decoder5(dec4)
+        dec6 = self.decoder6(dec5)
+        # Output layer
+        output = self.final_conv(dec6)
+        return output
 
 class ConvolutionalBlock(nn.Module):
     """