-
Notifications
You must be signed in to change notification settings - Fork 110
/
models.py
342 lines (268 loc) · 14.8 KB
/
models.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
import torch
from torch import nn
import torchvision
import math
class ConvolutionalBlock(nn.Module):
"""
A convolutional block, comprising convolutional, BN, activation layers.
"""
def __init__(self, in_channels, out_channels, kernel_size, stride=1, batch_norm=False, activation=None):
"""
:param in_channels: number of input channels
:param out_channels: number of output channe;s
:param kernel_size: kernel size
:param stride: stride
:param batch_norm: include a BN layer?
:param activation: Type of activation; None if none
"""
super(ConvolutionalBlock, self).__init__()
if activation is not None:
activation = activation.lower()
assert activation in {'prelu', 'leakyrelu', 'tanh'}
# A container that will hold the layers in this convolutional block
layers = list()
# A convolutional layer
layers.append(
nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride,
padding=kernel_size // 2))
# A batch normalization (BN) layer, if wanted
if batch_norm is True:
layers.append(nn.BatchNorm2d(num_features=out_channels))
# An activation layer, if wanted
if activation == 'prelu':
layers.append(nn.PReLU())
elif activation == 'leakyrelu':
layers.append(nn.LeakyReLU(0.2))
elif activation == 'tanh':
layers.append(nn.Tanh())
# Put together the convolutional block as a sequence of the layers in this container
self.conv_block = nn.Sequential(*layers)
def forward(self, input):
"""
Forward propagation.
:param input: input images, a tensor of size (N, in_channels, w, h)
:return: output images, a tensor of size (N, out_channels, w, h)
"""
output = self.conv_block(input) # (N, out_channels, w, h)
return output
class SubPixelConvolutionalBlock(nn.Module):
"""
A subpixel convolutional block, comprising convolutional, pixel-shuffle, and PReLU activation layers.
"""
def __init__(self, kernel_size=3, n_channels=64, scaling_factor=2):
"""
:param kernel_size: kernel size of the convolution
:param n_channels: number of input and output channels
:param scaling_factor: factor to scale input images by (along both dimensions)
"""
super(SubPixelConvolutionalBlock, self).__init__()
# A convolutional layer that increases the number of channels by scaling factor^2, followed by pixel shuffle and PReLU
self.conv = nn.Conv2d(in_channels=n_channels, out_channels=n_channels * (scaling_factor ** 2),
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.prelu = nn.PReLU()
def forward(self, input):
"""
Forward propagation.
:param input: input images, a tensor of size (N, n_channels, w, h)
:return: scaled output images, a tensor of size (N, n_channels, w * scaling factor, h * scaling factor)
"""
output = self.conv(input) # (N, n_channels * scaling factor^2, w, h)
output = self.pixel_shuffle(output) # (N, n_channels, w * scaling factor, h * scaling factor)
output = self.prelu(output) # (N, n_channels, w * scaling factor, h * scaling factor)
return output
class ResidualBlock(nn.Module):
"""
A residual block, comprising two convolutional blocks with a residual connection across them.
"""
def __init__(self, kernel_size=3, n_channels=64):
"""
:param kernel_size: kernel size
:param n_channels: number of input and output channels (same because the input must be added to the output)
"""
super(ResidualBlock, self).__init__()
# The first convolutional block
self.conv_block1 = ConvolutionalBlock(in_channels=n_channels, out_channels=n_channels, kernel_size=kernel_size,
batch_norm=True, activation='PReLu')
# The second convolutional block
self.conv_block2 = ConvolutionalBlock(in_channels=n_channels, out_channels=n_channels, kernel_size=kernel_size,
batch_norm=True, activation=None)
def forward(self, input):
"""
Forward propagation.
:param input: input images, a tensor of size (N, n_channels, w, h)
:return: output images, a tensor of size (N, n_channels, w, h)
"""
residual = input # (N, n_channels, w, h)
output = self.conv_block1(input) # (N, n_channels, w, h)
output = self.conv_block2(output) # (N, n_channels, w, h)
output = output + residual # (N, n_channels, w, h)
return output
class SRResNet(nn.Module):
"""
The SRResNet, as defined in the paper.
"""
def __init__(self, large_kernel_size=9, small_kernel_size=3, n_channels=64, n_blocks=16, scaling_factor=4):
"""
:param large_kernel_size: kernel size of the first and last convolutions which transform the inputs and outputs
:param small_kernel_size: kernel size of all convolutions in-between, i.e. those in the residual and subpixel convolutional blocks
:param n_channels: number of channels in-between, i.e. the input and output channels for the residual and subpixel convolutional blocks
:param n_blocks: number of residual blocks
:param scaling_factor: factor to scale input images by (along both dimensions) in the subpixel convolutional block
"""
super(SRResNet, self).__init__()
# Scaling factor must be 2, 4, or 8
scaling_factor = int(scaling_factor)
assert scaling_factor in {2, 4, 8}, "The scaling factor must be 2, 4, or 8!"
# The first convolutional block
self.conv_block1 = ConvolutionalBlock(in_channels=3, out_channels=n_channels, kernel_size=large_kernel_size,
batch_norm=False, activation='PReLu')
# A sequence of n_blocks residual blocks, each containing a skip-connection across the block
self.residual_blocks = nn.Sequential(
*[ResidualBlock(kernel_size=small_kernel_size, n_channels=n_channels) for i in range(n_blocks)])
# Another convolutional block
self.conv_block2 = ConvolutionalBlock(in_channels=n_channels, out_channels=n_channels,
kernel_size=small_kernel_size,
batch_norm=True, activation=None)
# Upscaling is done by sub-pixel convolution, with each such block upscaling by a factor of 2
n_subpixel_convolution_blocks = int(math.log2(scaling_factor))
self.subpixel_convolutional_blocks = nn.Sequential(
*[SubPixelConvolutionalBlock(kernel_size=small_kernel_size, n_channels=n_channels, scaling_factor=2) for i
in range(n_subpixel_convolution_blocks)])
# The last convolutional block
self.conv_block3 = ConvolutionalBlock(in_channels=n_channels, out_channels=3, kernel_size=large_kernel_size,
batch_norm=False, activation='Tanh')
def forward(self, lr_imgs):
"""
Forward prop.
:param lr_imgs: low-resolution input images, a tensor of size (N, 3, w, h)
:return: super-resolution output images, a tensor of size (N, 3, w * scaling factor, h * scaling factor)
"""
output = self.conv_block1(lr_imgs) # (N, 3, w, h)
residual = output # (N, n_channels, w, h)
output = self.residual_blocks(output) # (N, n_channels, w, h)
output = self.conv_block2(output) # (N, n_channels, w, h)
output = output + residual # (N, n_channels, w, h)
output = self.subpixel_convolutional_blocks(output) # (N, n_channels, w * scaling factor, h * scaling factor)
sr_imgs = self.conv_block3(output) # (N, 3, w * scaling factor, h * scaling factor)
return sr_imgs
class Generator(nn.Module):
"""
The generator in the SRGAN, as defined in the paper. Architecture identical to the SRResNet.
"""
def __init__(self, large_kernel_size=9, small_kernel_size=3, n_channels=64, n_blocks=16, scaling_factor=4):
"""
:param large_kernel_size: kernel size of the first and last convolutions which transform the inputs and outputs
:param small_kernel_size: kernel size of all convolutions in-between, i.e. those in the residual and subpixel convolutional blocks
:param n_channels: number of channels in-between, i.e. the input and output channels for the residual and subpixel convolutional blocks
:param n_blocks: number of residual blocks
:param scaling_factor: factor to scale input images by (along both dimensions) in the subpixel convolutional block
"""
super(Generator, self).__init__()
# The generator is simply an SRResNet, as above
self.net = SRResNet(large_kernel_size=large_kernel_size, small_kernel_size=small_kernel_size,
n_channels=n_channels, n_blocks=n_blocks, scaling_factor=scaling_factor)
def initialize_with_srresnet(self, srresnet_checkpoint):
"""
Initialize with weights from a trained SRResNet.
:param srresnet_checkpoint: checkpoint filepath
"""
srresnet = torch.load(srresnet_checkpoint)['model']
self.net.load_state_dict(srresnet.state_dict())
print("\nLoaded weights from pre-trained SRResNet.\n")
def forward(self, lr_imgs):
"""
Forward prop.
:param lr_imgs: low-resolution input images, a tensor of size (N, 3, w, h)
:return: super-resolution output images, a tensor of size (N, 3, w * scaling factor, h * scaling factor)
"""
sr_imgs = self.net(lr_imgs) # (N, n_channels, w * scaling factor, h * scaling factor)
return sr_imgs
class Discriminator(nn.Module):
"""
The discriminator in the SRGAN, as defined in the paper.
"""
def __init__(self, kernel_size=3, n_channels=64, n_blocks=8, fc_size=1024):
"""
:param kernel_size: kernel size in all convolutional blocks
:param n_channels: number of output channels in the first convolutional block, after which it is doubled in every 2nd block thereafter
:param n_blocks: number of convolutional blocks
:param fc_size: size of the first fully connected layer
"""
super(Discriminator, self).__init__()
in_channels = 3
# A series of convolutional blocks
# The first, third, fifth (and so on) convolutional blocks increase the number of channels but retain image size
# The second, fourth, sixth (and so on) convolutional blocks retain the same number of channels but halve image size
# The first convolutional block is unique because it does not employ batch normalization
conv_blocks = list()
for i in range(n_blocks):
out_channels = (n_channels if i is 0 else in_channels * 2) if i % 2 is 0 else in_channels
conv_blocks.append(
ConvolutionalBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,
stride=1 if i % 2 is 0 else 2, batch_norm=i is not 0, activation='LeakyReLu'))
in_channels = out_channels
self.conv_blocks = nn.Sequential(*conv_blocks)
# An adaptive pool layer that resizes it to a standard size
# For the default input size of 96 and 8 convolutional blocks, this will have no effect
self.adaptive_pool = nn.AdaptiveAvgPool2d((6, 6))
self.fc1 = nn.Linear(out_channels * 6 * 6, fc_size)
self.leaky_relu = nn.LeakyReLU(0.2)
self.fc2 = nn.Linear(1024, 1)
# Don't need a sigmoid layer because the sigmoid operation is performed by PyTorch's nn.BCEWithLogitsLoss()
def forward(self, imgs):
"""
Forward propagation.
:param imgs: high-resolution or super-resolution images which must be classified as such, a tensor of size (N, 3, w * scaling factor, h * scaling factor)
:return: a score (logit) for whether it is a high-resolution image, a tensor of size (N)
"""
batch_size = imgs.size(0)
output = self.conv_blocks(imgs)
output = self.adaptive_pool(output)
output = self.fc1(output.view(batch_size, -1))
output = self.leaky_relu(output)
logit = self.fc2(output)
return logit
class TruncatedVGG19(nn.Module):
"""
A truncated VGG19 network, such that its output is the 'feature map obtained by the j-th convolution (after activation)
before the i-th maxpooling layer within the VGG19 network', as defined in the paper.
Used to calculate the MSE loss in this VGG feature-space, i.e. the VGG loss.
"""
def __init__(self, i, j):
"""
:param i: the index i in the definition above
:param j: the index j in the definition above
"""
super(TruncatedVGG19, self).__init__()
# Load the pre-trained VGG19 available in torchvision
vgg19 = torchvision.models.vgg19(pretrained=True)
maxpool_counter = 0
conv_counter = 0
truncate_at = 0
# Iterate through the convolutional section ("features") of the VGG19
for layer in vgg19.features.children():
truncate_at += 1
# Count the number of maxpool layers and the convolutional layers after each maxpool
if isinstance(layer, nn.Conv2d):
conv_counter += 1
if isinstance(layer, nn.MaxPool2d):
maxpool_counter += 1
conv_counter = 0
# Break if we reach the jth convolution after the (i - 1)th maxpool
if maxpool_counter == i - 1 and conv_counter == j:
break
# Check if conditions were satisfied
assert maxpool_counter == i - 1 and conv_counter == j, "One or both of i=%d and j=%d are not valid choices for the VGG19!" % (
i, j)
# Truncate to the jth convolution (+ activation) before the ith maxpool layer
self.truncated_vgg19 = nn.Sequential(*list(vgg19.features.children())[:truncate_at + 1])
def forward(self, input):
"""
Forward propagation
:param input: high-resolution or super-resolution images, a tensor of size (N, 3, w * scaling factor, h * scaling factor)
:return: the specified VGG19 feature map, a tensor of size (N, feature_map_channels, feature_map_w, feature_map_h)
"""
output = self.truncated_vgg19(input) # (N, feature_map_channels, feature_map_w, feature_map_h)
return output