Update realesrgan/archs/rrdbnet_arch.py

realesrgan/archs/rrdbnet_arch.py

@@ -1,52 +1,61 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-class ResidualDenseBlock(nn.Module):
-    def __init__(self, nf=64, gc=32):
-        super().__init__()
-        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1)
-        self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1)
-        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1)
-        self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1)
-        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1)
-        self.lrelu = nn.LeakyReLU(0.2, inplace=True)
-
-    def forward(self, x):
-        x1 = self.lrelu(self.conv1(x))
-        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
-        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
-        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
-        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
-        return x + 0.2 * x5
-
-class RRDB(nn.Module):
-    def __init__(self, nf, gc=32):
-        super().__init__()
-        self.rdb1 = ResidualDenseBlock(nf, gc)
-        self.rdb2 = ResidualDenseBlock(nf, gc)
-        self.rdb3 = ResidualDenseBlock(nf, gc)
-
-    def forward(self, x):
-        return x + 0.2 * self.rdb3(self.rdb2(self.rdb1(x)))
-
-class RRDBNet(nn.Module):
-    def __init__(self, in_nc, out_nc, nf, nb, gc=32):
-        super().__init__()
-        self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1)
-        self.RRDB_trunk = nn.Sequential(*[RRDB(nf, gc) for _ in range(nb)])
-        self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1)
-        self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1)
-        self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1)
-        self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1)
-        self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1)
-        self.lrelu = nn.LeakyReLU(0.2, inplace=True)
-
-    def forward(self, x):
-        fea = self.conv_first(x)
-        trunk = self.trunk_conv(self.RRDB_trunk(fea))
-        fea = fea + trunk
-        fea = self.lrelu(self.upconv1(torch.nn.functional.interpolate(fea, scale_factor=2, mode='nearest')))
-        fea = self.lrelu(self.upconv2(torch.nn.functional.interpolate(fea, scale_factor=2, mode='nearest')))
-        out = self.conv_last(self.lrelu(self.HRconv(fea)))
-        return out
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ResidualDenseBlock(nn.Module):
+    """Residual Dense Block used in RRDB"""
+
+    def __init__(self, nf=64, gc=32):
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1)
+        self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1)
+        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1)
+        self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1)
+        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1)
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x):
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        return x + 0.2 * x5
+
+
+class RRDB(nn.Module):
+    """Residual in Residual Dense Block"""
+
+    def __init__(self, nf, gc=32):
+        super(RRDB, self).__init__()
+        self.rdb1 = ResidualDenseBlock(nf, gc)
+        self.rdb2 = ResidualDenseBlock(nf, gc)
+        self.rdb3 = ResidualDenseBlock(nf, gc)
+
+    def forward(self, x):
+        return x + 0.2 * self.rdb3(self.rdb2(self.rdb1(x)))
+
+
+class RRDBNet(nn.Module):
+    def __init__(self, in_nc, out_nc, nf, nb, gc=32):
+        super(RRDBNet, self).__init__()
+        self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1)
+        self.body = nn.Sequential(*[RRDB(nf, gc) for _ in range(nb)])
+        self.conv_body = nn.Conv2d(nf, nf, 3, 1, 1)
+
+        # upsampling
+        self.conv_up1 = nn.Conv2d(nf, nf, 3, 1, 1)
+        self.conv_up2 = nn.Conv2d(nf, nf, 3, 1, 1)
+        self.conv_hr = nn.Conv2d(nf, nf, 3, 1, 1)
+        self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1)
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x):
+        fea = self.conv_first(x)
+        trunk = self.conv_body(self.body(fea))
+        fea = fea + trunk
+        fea = self.lrelu(self.conv_up1(F.interpolate(fea, scale_factor=2, mode='nearest')))
+        fea = self.lrelu(self.conv_up2(F.interpolate(fea, scale_factor=2, mode='nearest')))
+        out = self.conv_last(self.lrelu(self.conv_hr(fea)))
+        return out