-
Notifications
You must be signed in to change notification settings - Fork 0
/
models.py
185 lines (149 loc) · 8.16 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
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