Delete model/SRMNet

model/SRMNet

@@ -1,227 +0,0 @@
-import torch
-import torch.nn as nn
-
-##---------- Basic Layers ----------
-def conv3x3(in_chn, out_chn, bias=True):
-    layer = nn.Conv2d(in_chn, out_chn, kernel_size=3, stride=1, padding=1, bias=bias)
-    return layer
-
-def conv(in_channels, out_channels, kernel_size, bias=False, stride=1):
-    return nn.Conv2d(
-        in_channels, out_channels, kernel_size,
-        padding=(kernel_size // 2), bias=bias, stride=stride)
-
-def bili_resize(factor):
-    return nn.Upsample(scale_factor=factor, mode='bilinear', align_corners=False)
-
-##---------- Basic Blocks ----------
-class UNetConvBlock(nn.Module):
-    def __init__(self, in_size, out_size, downsample):
-        super(UNetConvBlock, self).__init__()
-        self.downsample = downsample
-        self.block = SK_RDB(in_channels=in_size, growth_rate=out_size, num_layers=3)
-        if downsample:
-            self.downsample = PS_down(out_size, out_size, downscale=2)
-
-    def forward(self, x):
-        out = self.block(x)
-        if self.downsample:
-            out_down = self.downsample(out)
-            return out_down, out
-        else:
-            return out
-
-class UNetUpBlock(nn.Module):
-    def __init__(self, in_size, out_size):
-        super(UNetUpBlock, self).__init__()
-        # self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2, bias=True)
-        self.up = PS_up(in_size, out_size, upscale=2)
-        self.conv_block = UNetConvBlock(in_size, out_size, False)
-
-    def forward(self, x, bridge):
-        up = self.up(x)
-        out = torch.cat([up, bridge], dim=1)
-        out = self.conv_block(out)
-        return out
-
-##---------- Resizing Modules (Pixel(Un)Shuffle) ----------
-class PS_down(nn.Module):
-    def __init__(self, in_size, out_size, downscale):
-        super(PS_down, self).__init__()
-        self.UnPS = nn.PixelUnshuffle(downscale)
-        self.conv1 = nn.Conv2d((downscale**2) * in_size, out_size, 1, 1, 0)
-
-    def forward(self, x):
-        x = self.UnPS(x)  # h/2, w/2, 4*c
-        x = self.conv1(x)
-        return x
-
-class PS_up(nn.Module):
-    def __init__(self, in_size, out_size, upscale):
-        super(PS_up, self).__init__()
-
-        self.PS = nn.PixelShuffle(upscale)
-        self.conv1 = nn.Conv2d(in_size//(upscale**2), out_size, 1, 1, 0)
-
-    def forward(self, x):
-        x = self.PS(x)  # h/2, w/2, 4*c
-        x = self.conv1(x)
-        return x
-
-##---------- Selective Kernel Feature Fusion (SKFF) ----------
-class SKFF(nn.Module):
-    def __init__(self, in_channels, height=3, reduction=8, bias=False):
-        super(SKFF, self).__init__()
-
-        self.height = height
-        d = max(int(in_channels / reduction), 4)
-
-        self.avg_pool = nn.AdaptiveAvgPool2d(1)
-        self.conv_du = nn.Sequential(nn.Conv2d(in_channels, d, 1, padding=0, bias=bias), nn.PReLU())
-
-        self.fcs = nn.ModuleList([])
-        for i in range(self.height):
-            self.fcs.append(nn.Conv2d(d, in_channels, kernel_size=1, stride=1, bias=bias))
-
-        self.softmax = nn.Softmax(dim=1)
-
-    def forward(self, inp_feats):
-        batch_size, n_feats, H, W = inp_feats[1].shape
-
-        inp_feats = torch.cat(inp_feats, dim=1)
-        inp_feats = inp_feats.view(batch_size, self.height, n_feats, inp_feats.shape[2], inp_feats.shape[3])
-
-        feats_U = torch.sum(inp_feats, dim=1)
-        feats_S = self.avg_pool(feats_U)
-        feats_Z = self.conv_du(feats_S)
-
-        attention_vectors = [fc(feats_Z) for fc in self.fcs]
-        attention_vectors = torch.cat(attention_vectors, dim=1)
-        attention_vectors = attention_vectors.view(batch_size, self.height, n_feats, 1, 1)
-
-        attention_vectors = self.softmax(attention_vectors)
-        feats_V = torch.sum(inp_feats * attention_vectors, dim=1)
-
-        return feats_V
-
-##---------- Dense Block ----------
-class DenseLayer(nn.Module):
-    def __init__(self, in_channels, out_channels, I):
-        super(DenseLayer, self).__init__()
-        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=3 // 2)
-        self.relu = nn.ReLU(inplace=True)
-        self.sk = SKFF(out_channels, height=2, reduction=8, bias=False)
-
-    def forward(self, x):
-        x1 = self.relu(self.conv(x))
-        # output = torch.cat([x, x1], 1) # -> RDB
-        output = self.sk((x, x1))
-        return output
-
-##---------- Selective Kernel Residual Dense Block (SK-RDB) ----------
-class SK_RDB(nn.Module):
-    def __init__(self, in_channels, growth_rate, num_layers):
-        super(SK_RDB, self).__init__()
-        self.identity = nn.Conv2d(in_channels, growth_rate, 1, 1, 0)
-        self.layers = nn.Sequential(
-            *[DenseLayer(in_channels, in_channels, I=i) for i in range(num_layers)]
-        )
-        self.lff = nn.Conv2d(in_channels, growth_rate, kernel_size=1)
-
-    def forward(self, x):
-        res = self.identity(x)
-        x = self.layers(x)
-        x = self.lff(x)
-        return res + x
-
-##---------- testNet ----------
-class SRMNet(nn.Module):
-    def __init__(self, in_chn=3, wf=96, depth=4):
-        super(SRMNet, self).__init__()
-        self.depth = depth
-        self.down_path = nn.ModuleList()
-        self.bili_down = bili_resize(0.5)
-        self.conv_01 = nn.Conv2d(in_chn, wf, 3, 1, 1)
-
-        # encoder of UNet
-        prev_channels = 0
-        for i in range(depth):  # 0,1,2,3
-            downsample = True if (i + 1) < depth else False
-            self.down_path.append(UNetConvBlock(prev_channels + wf, (2 ** i) * wf, downsample))
-            prev_channels = (2 ** i) * wf
-
-        # decoder of UNet
-        self.up_path = nn.ModuleList()
-        self.skip_conv = nn.ModuleList()
-        self.conv_up = nn.ModuleList()
-        self.bottom_conv = nn.Conv2d(prev_channels, wf, 3, 1, 1)
-        self.bottom_up = bili_resize(2 ** (depth-1))
-
-        for i in reversed(range(depth - 1)):
-            self.up_path.append(UNetUpBlock(prev_channels, (2 ** i) * wf))
-            self.skip_conv.append(nn.Conv2d((2 ** i) * wf, (2 ** i) * wf, 3, 1, 1))
-            self.conv_up.append(nn.Sequential(*[nn.Conv2d((2 ** i) * wf, wf, 3, 1, 1), bili_resize(2 ** i)]))
-            prev_channels = (2 ** i) * wf
-
-        self.final_ff = SKFF(in_channels=wf, height=depth)
-        self.last = conv3x3(prev_channels, in_chn, bias=True)
-
-    def forward(self, x):
-        img = x
-        scale_img = img
-
-        ##### shallow conv #####
-        x1 = self.conv_01(img)
-        encs = []
-        ######## UNet ########
-        # Down-path (Encoder)
-        for i, down in enumerate(self.down_path):
-            if i == 0:
-                x1, x1_up = down(x1)
-                encs.append(x1_up)
-            elif (i + 1) < self.depth:
-                scale_img = self.bili_down(scale_img)
-                left_bar = self.conv_01(scale_img)
-                x1 = torch.cat([x1, left_bar], dim=1)
-                x1, x1_up = down(x1)
-                encs.append(x1_up)
-            else:
-                scale_img = self.bili_down(scale_img)
-                left_bar = self.conv_01(scale_img)
-                x1 = torch.cat([x1, left_bar], dim=1)
-                x1 = down(x1)
-
-        # Up-path (Decoder)
-        ms_result = [self.bottom_up(self.bottom_conv(x1))]
-        for i, up in enumerate(self.up_path):
-            x1 = up(x1, self.skip_conv[i](encs[-i - 1]))
-            ms_result.append(self.conv_up[i](x1))
-
-        # Multi-scale selective feature fusion
-        msff_result = self.final_ff(ms_result)
-
-        ##### Reconstruct #####
-        out_1 = self.last(msff_result) + img
-
-        return out_1
-
-
-if __name__ == "__main__":
-    from thop import profile
-
-    input = torch.ones(1, 3, 256, 256, dtype=torch.float, requires_grad=False)
-    model = SRMNet(in_chn=3, wf=96, depth=4)
-    out = model(input)
-    flops, params = profile(model, inputs=(input,))
-    total = sum(p.numel() for p in model.parameters())
-
-    # RDBlayer = SK_RDB(in_channels=64, growth_rate=64, num_layers=3)
-    # print(RDBlayer)
-    # out = RDBlayer(input)
-    # flops, params = profile(RDBlayer, inputs=(input,))
-
-    print('input shape:', input.shape)
-    print('output shape', out.shape)
-    print("-----------------------------------")
-    print("Total params: %.4f M" % (total / 1e6))
-    print("Total params: %.4f G" % (flops / 1e9))