1 files changed, 135 insertions, 0 deletions
diff --git a/r_basicsr/archs/rcan_arch.py b/r_basicsr/archs/rcan_arch.py
new file mode 100644
index 0000000..78f917e
--- /dev/null
+++ b/r_basicsr/archs/rcan_arch.py
@@ -0,0 +1,135 @@
+import torch
+from torch import nn as nn
+
+from r_basicsr.utils.registry import ARCH_REGISTRY
+from .arch_util import Upsample, make_layer
+
+
+class ChannelAttention(nn.Module):
+    """Channel attention used in RCAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        squeeze_factor (int): Channel squeeze factor. Default: 16.
+    """
+
+    def __init__(self, num_feat, squeeze_factor=16):
+        super(ChannelAttention, self).__init__()
+        self.attention = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1), nn.Conv2d(num_feat, num_feat // squeeze_factor, 1, padding=0),
+            nn.ReLU(inplace=True), nn.Conv2d(num_feat // squeeze_factor, num_feat, 1, padding=0), nn.Sigmoid())
+
+    def forward(self, x):
+        y = self.attention(x)
+        return x * y
+
+
+class RCAB(nn.Module):
+    """Residual Channel Attention Block (RCAB) used in RCAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        squeeze_factor (int): Channel squeeze factor. Default: 16.
+        res_scale (float): Scale the residual. Default: 1.
+    """
+
+    def __init__(self, num_feat, squeeze_factor=16, res_scale=1):
+        super(RCAB, self).__init__()
+        self.res_scale = res_scale
+
+        self.rcab = nn.Sequential(
+            nn.Conv2d(num_feat, num_feat, 3, 1, 1), nn.ReLU(True), nn.Conv2d(num_feat, num_feat, 3, 1, 1),
+            ChannelAttention(num_feat, squeeze_factor))
+
+    def forward(self, x):
+        res = self.rcab(x) * self.res_scale
+        return res + x
+
+
+class ResidualGroup(nn.Module):
+    """Residual Group of RCAB.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_block (int): Block number in the body network.
+        squeeze_factor (int): Channel squeeze factor. Default: 16.
+        res_scale (float): Scale the residual. Default: 1.
+    """
+
+    def __init__(self, num_feat, num_block, squeeze_factor=16, res_scale=1):
+        super(ResidualGroup, self).__init__()
+
+        self.residual_group = make_layer(
+            RCAB, num_block, num_feat=num_feat, squeeze_factor=squeeze_factor, res_scale=res_scale)
+        self.conv = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+
+    def forward(self, x):
+        res = self.conv(self.residual_group(x))
+        return res + x
+
+
+@ARCH_REGISTRY.register()
+class RCAN(nn.Module):
+    """Residual Channel Attention Networks.
+
+    Paper: Image Super-Resolution Using Very Deep Residual Channel Attention
+        Networks
+    Ref git repo: https://github.com/yulunzhang/RCAN.
+
+    Args:
+        num_in_ch (int): Channel number of inputs.
+        num_out_ch (int): Channel number of outputs.
+        num_feat (int): Channel number of intermediate features.
+            Default: 64.
+        num_group (int): Number of ResidualGroup. Default: 10.
+        num_block (int): Number of RCAB in ResidualGroup. Default: 16.
+        squeeze_factor (int): Channel squeeze factor. Default: 16.
+        upscale (int): Upsampling factor. Support 2^n and 3.
+            Default: 4.
+        res_scale (float): Used to scale the residual in residual block.
+            Default: 1.
+        img_range (float): Image range. Default: 255.
+        rgb_mean (tuple[float]): Image mean in RGB orders.
+            Default: (0.4488, 0.4371, 0.4040), calculated from DIV2K dataset.
+    """
+
+    def __init__(self,
+                 num_in_ch,
+                 num_out_ch,
+                 num_feat=64,
+                 num_group=10,
+                 num_block=16,
+                 squeeze_factor=16,
+                 upscale=4,
+                 res_scale=1,
+                 img_range=255.,
+                 rgb_mean=(0.4488, 0.4371, 0.4040)):
+        super(RCAN, self).__init__()
+
+        self.img_range = img_range
+        self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
+
+        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
+        self.body = make_layer(
+            ResidualGroup,
+            num_group,
+            num_feat=num_feat,
+            num_block=num_block,
+            squeeze_factor=squeeze_factor,
+            res_scale=res_scale)
+        self.conv_after_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.upsample = Upsample(upscale, num_feat)
+        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+
+    def forward(self, x):
+        self.mean = self.mean.type_as(x)
+
+        x = (x - self.mean) * self.img_range
+        x = self.conv_first(x)
+        res = self.conv_after_body(self.body(x))
+        res += x
+
+        x = self.conv_last(self.upsample(res))
+        x = x / self.img_range + self.mean
+
+        return x