52Hz commited on
Commit
c9e5229
1 Parent(s): 1889802

Create SRMNet.py

Browse files
Files changed (1) hide show
  1. model/SRMNet.py +227 -0
model/SRMNet.py ADDED
@@ -0,0 +1,227 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+
4
+ ##---------- Basic Layers ----------
5
+ def conv3x3(in_chn, out_chn, bias=True):
6
+ layer = nn.Conv2d(in_chn, out_chn, kernel_size=3, stride=1, padding=1, bias=bias)
7
+ return layer
8
+
9
+ def conv(in_channels, out_channels, kernel_size, bias=False, stride=1):
10
+ return nn.Conv2d(
11
+ in_channels, out_channels, kernel_size,
12
+ padding=(kernel_size // 2), bias=bias, stride=stride)
13
+
14
+ def bili_resize(factor):
15
+ return nn.Upsample(scale_factor=factor, mode='bilinear', align_corners=False)
16
+
17
+ ##---------- Basic Blocks ----------
18
+ class UNetConvBlock(nn.Module):
19
+ def __init__(self, in_size, out_size, downsample):
20
+ super(UNetConvBlock, self).__init__()
21
+ self.downsample = downsample
22
+ self.block = SK_RDB(in_channels=in_size, growth_rate=out_size, num_layers=3)
23
+ if downsample:
24
+ self.downsample = PS_down(out_size, out_size, downscale=2)
25
+
26
+ def forward(self, x):
27
+ out = self.block(x)
28
+ if self.downsample:
29
+ out_down = self.downsample(out)
30
+ return out_down, out
31
+ else:
32
+ return out
33
+
34
+ class UNetUpBlock(nn.Module):
35
+ def __init__(self, in_size, out_size):
36
+ super(UNetUpBlock, self).__init__()
37
+ # self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2, bias=True)
38
+ self.up = PS_up(in_size, out_size, upscale=2)
39
+ self.conv_block = UNetConvBlock(in_size, out_size, False)
40
+
41
+ def forward(self, x, bridge):
42
+ up = self.up(x)
43
+ out = torch.cat([up, bridge], dim=1)
44
+ out = self.conv_block(out)
45
+ return out
46
+
47
+ ##---------- Resizing Modules (Pixel(Un)Shuffle) ----------
48
+ class PS_down(nn.Module):
49
+ def __init__(self, in_size, out_size, downscale):
50
+ super(PS_down, self).__init__()
51
+ self.UnPS = nn.PixelUnshuffle(downscale)
52
+ self.conv1 = nn.Conv2d((downscale**2) * in_size, out_size, 1, 1, 0)
53
+
54
+ def forward(self, x):
55
+ x = self.UnPS(x) # h/2, w/2, 4*c
56
+ x = self.conv1(x)
57
+ return x
58
+
59
+ class PS_up(nn.Module):
60
+ def __init__(self, in_size, out_size, upscale):
61
+ super(PS_up, self).__init__()
62
+
63
+ self.PS = nn.PixelShuffle(upscale)
64
+ self.conv1 = nn.Conv2d(in_size//(upscale**2), out_size, 1, 1, 0)
65
+
66
+ def forward(self, x):
67
+ x = self.PS(x) # h/2, w/2, 4*c
68
+ x = self.conv1(x)
69
+ return x
70
+
71
+ ##---------- Selective Kernel Feature Fusion (SKFF) ----------
72
+ class SKFF(nn.Module):
73
+ def __init__(self, in_channels, height=3, reduction=8, bias=False):
74
+ super(SKFF, self).__init__()
75
+
76
+ self.height = height
77
+ d = max(int(in_channels / reduction), 4)
78
+
79
+ self.avg_pool = nn.AdaptiveAvgPool2d(1)
80
+ self.conv_du = nn.Sequential(nn.Conv2d(in_channels, d, 1, padding=0, bias=bias), nn.PReLU())
81
+
82
+ self.fcs = nn.ModuleList([])
83
+ for i in range(self.height):
84
+ self.fcs.append(nn.Conv2d(d, in_channels, kernel_size=1, stride=1, bias=bias))
85
+
86
+ self.softmax = nn.Softmax(dim=1)
87
+
88
+ def forward(self, inp_feats):
89
+ batch_size, n_feats, H, W = inp_feats[1].shape
90
+
91
+ inp_feats = torch.cat(inp_feats, dim=1)
92
+ inp_feats = inp_feats.view(batch_size, self.height, n_feats, inp_feats.shape[2], inp_feats.shape[3])
93
+
94
+ feats_U = torch.sum(inp_feats, dim=1)
95
+ feats_S = self.avg_pool(feats_U)
96
+ feats_Z = self.conv_du(feats_S)
97
+
98
+ attention_vectors = [fc(feats_Z) for fc in self.fcs]
99
+ attention_vectors = torch.cat(attention_vectors, dim=1)
100
+ attention_vectors = attention_vectors.view(batch_size, self.height, n_feats, 1, 1)
101
+
102
+ attention_vectors = self.softmax(attention_vectors)
103
+ feats_V = torch.sum(inp_feats * attention_vectors, dim=1)
104
+
105
+ return feats_V
106
+
107
+ ##---------- Dense Block ----------
108
+ class DenseLayer(nn.Module):
109
+ def __init__(self, in_channels, out_channels, I):
110
+ super(DenseLayer, self).__init__()
111
+ self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=3 // 2)
112
+ self.relu = nn.ReLU(inplace=True)
113
+ self.sk = SKFF(out_channels, height=2, reduction=8, bias=False)
114
+
115
+ def forward(self, x):
116
+ x1 = self.relu(self.conv(x))
117
+ # output = torch.cat([x, x1], 1) # -> RDB
118
+ output = self.sk((x, x1))
119
+ return output
120
+
121
+ ##---------- Selective Kernel Residual Dense Block (SK-RDB) ----------
122
+ class SK_RDB(nn.Module):
123
+ def __init__(self, in_channels, growth_rate, num_layers):
124
+ super(SK_RDB, self).__init__()
125
+ self.identity = nn.Conv2d(in_channels, growth_rate, 1, 1, 0)
126
+ self.layers = nn.Sequential(
127
+ *[DenseLayer(in_channels, in_channels, I=i) for i in range(num_layers)]
128
+ )
129
+ self.lff = nn.Conv2d(in_channels, growth_rate, kernel_size=1)
130
+
131
+ def forward(self, x):
132
+ res = self.identity(x)
133
+ x = self.layers(x)
134
+ x = self.lff(x)
135
+ return res + x
136
+
137
+ ##---------- testNet ----------
138
+ class SRMNet(nn.Module):
139
+ def __init__(self, in_chn=3, wf=96, depth=4):
140
+ super(SRMNet, self).__init__()
141
+ self.depth = depth
142
+ self.down_path = nn.ModuleList()
143
+ self.bili_down = bili_resize(0.5)
144
+ self.conv_01 = nn.Conv2d(in_chn, wf, 3, 1, 1)
145
+
146
+ # encoder of UNet
147
+ prev_channels = 0
148
+ for i in range(depth): # 0,1,2,3
149
+ downsample = True if (i + 1) < depth else False
150
+ self.down_path.append(UNetConvBlock(prev_channels + wf, (2 ** i) * wf, downsample))
151
+ prev_channels = (2 ** i) * wf
152
+
153
+ # decoder of UNet
154
+ self.up_path = nn.ModuleList()
155
+ self.skip_conv = nn.ModuleList()
156
+ self.conv_up = nn.ModuleList()
157
+ self.bottom_conv = nn.Conv2d(prev_channels, wf, 3, 1, 1)
158
+ self.bottom_up = bili_resize(2 ** (depth-1))
159
+
160
+ for i in reversed(range(depth - 1)):
161
+ self.up_path.append(UNetUpBlock(prev_channels, (2 ** i) * wf))
162
+ self.skip_conv.append(nn.Conv2d((2 ** i) * wf, (2 ** i) * wf, 3, 1, 1))
163
+ self.conv_up.append(nn.Sequential(*[nn.Conv2d((2 ** i) * wf, wf, 3, 1, 1), bili_resize(2 ** i)]))
164
+ prev_channels = (2 ** i) * wf
165
+
166
+ self.final_ff = SKFF(in_channels=wf, height=depth)
167
+ self.last = conv3x3(prev_channels, in_chn, bias=True)
168
+
169
+ def forward(self, x):
170
+ img = x
171
+ scale_img = img
172
+
173
+ ##### shallow conv #####
174
+ x1 = self.conv_01(img)
175
+ encs = []
176
+ ######## UNet ########
177
+ # Down-path (Encoder)
178
+ for i, down in enumerate(self.down_path):
179
+ if i == 0:
180
+ x1, x1_up = down(x1)
181
+ encs.append(x1_up)
182
+ elif (i + 1) < self.depth:
183
+ scale_img = self.bili_down(scale_img)
184
+ left_bar = self.conv_01(scale_img)
185
+ x1 = torch.cat([x1, left_bar], dim=1)
186
+ x1, x1_up = down(x1)
187
+ encs.append(x1_up)
188
+ else:
189
+ scale_img = self.bili_down(scale_img)
190
+ left_bar = self.conv_01(scale_img)
191
+ x1 = torch.cat([x1, left_bar], dim=1)
192
+ x1 = down(x1)
193
+
194
+ # Up-path (Decoder)
195
+ ms_result = [self.bottom_up(self.bottom_conv(x1))]
196
+ for i, up in enumerate(self.up_path):
197
+ x1 = up(x1, self.skip_conv[i](encs[-i - 1]))
198
+ ms_result.append(self.conv_up[i](x1))
199
+
200
+ # Multi-scale selective feature fusion
201
+ msff_result = self.final_ff(ms_result)
202
+
203
+ ##### Reconstruct #####
204
+ out_1 = self.last(msff_result) + img
205
+
206
+ return out_1
207
+
208
+
209
+ if __name__ == "__main__":
210
+ from thop import profile
211
+
212
+ input = torch.ones(1, 3, 256, 256, dtype=torch.float, requires_grad=False)
213
+ model = SRMNet(in_chn=3, wf=96, depth=4)
214
+ out = model(input)
215
+ flops, params = profile(model, inputs=(input,))
216
+ total = sum(p.numel() for p in model.parameters())
217
+
218
+ # RDBlayer = SK_RDB(in_channels=64, growth_rate=64, num_layers=3)
219
+ # print(RDBlayer)
220
+ # out = RDBlayer(input)
221
+ # flops, params = profile(RDBlayer, inputs=(input,))
222
+
223
+ print('input shape:', input.shape)
224
+ print('output shape', out.shape)
225
+ print("-----------------------------------")
226
+ print("Total params: %.4f M" % (total / 1e6))
227
+ print("Total params: %.4f G" % (flops / 1e9))