def __init__(self,
in_channels,
num_layers,
num_filters,
kernel_size,
relu_leakiness,
use_bn,
downsample,
use_act=True,
padding='zero'):
super(ConvEncodeUnit, self).__init__()
self.downsample = downsample
use_bias = not use_bn
modules = []
for i in range(num_layers):
modules += [
get_same_padding_layer(kernel_size, stride=1, mode=padding),
nn.Conv2d(in_channels,
num_filters,
kernel_size=kernel_size,
stride=1,
bias=use_bias)
]
in_channels = num_filters
if use_bn:
modules += [nn.BatchNorm2d(in_channels)]
if use_act:
modules += [nn.LeakyReLU(relu_leakiness, inplace=True)]
self.encode = nn.Sequential(*modules)
if downsample:
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
def __init__(self,
num_convs,
num_filters,
kernel_size,
relu_leakiness,
dilations,
padding='zero',
num_inputs=2,
num_outputs=2,
final_act=False):
super(ConvBlock, self).__init__()
in_channels = num_inputs
modules = []
for i in range(num_convs - 1):
modules += [
get_same_padding_layer(kernel_size,
stride=1,
mode=padding,
dilation=dilations[i]),
nn.Conv2d(in_channels,
num_filters,
kernel_size=kernel_size,
stride=1,
bias=True,
dilation=dilations[i]),
nn.LeakyReLU(relu_leakiness, inplace=True)
]
in_channels = num_filters
modules += [
get_same_padding_layer(kernel_size,
stride=1,
mode=padding,
dilation=dilations[-1]),
nn.Conv2d(in_channels,
num_outputs,
kernel_size=kernel_size,
stride=1,
bias=True,
dilation=dilations[-1])
]
if final_act:
modules.append(nn.LeakyReLU(relu_leakiness, inplace=True))
self.layers = nn.Sequential(*modules)
def __init__(self, in_channels, num_filters, kernel_size, use_norm_layers,
norm_layer, act_fn, relu_leakiness=None, padding='zero'):
"""Builds a residual block
The implementation follows the SRGAN paper, which uses a residual block
variant which uses no activation after the addition. This design is based
on Sam Gross & Michael Wilber: "Training and investigating Residual Nets"
(see http://torch.ch/blog/2016/02/04/resnets.html)
Parameters
----------
in_channels : int
Number of input channels
num_filters : int
Number of convolutional filters to use
kernel_size : int
Size of convolution kernel
use_norm_layers : bool
If true, uses normalization layers after the convolution layers
norm_layer : string
Normalization layer to use. `batch` for batch normalization or `instance`
for instance normalization
act_fn : string
Activation function to use. Either `relu`, `prelu` (default), or `lrelu`
relu_leakiness : float
If using lrelu, leakiness of the relus, if using prelu, initial value
for prelu parameters
padding : string
Type of padding to use. Either `zero`, `reflection`, or `replication`
"""
super(ResBlock, self).__init__()
use_bias = need_bias(use_norm_layers, norm_layer)
modules = [get_same_padding_layer(kernel_size, stride=1, mode=padding),
nn.Conv2d(in_channels, num_filters, kernel_size=kernel_size,
stride=1, bias=use_bias)]
if use_norm_layers:
modules.append(get_normalization_layer(norm_layer, num_filters))
modules += [get_activation_fn(act_fn, relu_leakiness, num_filters),
get_same_padding_layer(kernel_size, stride=1, mode=padding),
nn.Conv2d(num_filters, num_filters, kernel_size=kernel_size,
stride=1, bias=use_bias)]
if use_norm_layers:
modules.append(get_normalization_layer(norm_layer, num_filters))
self.block = nn.Sequential(*modules)
def __init__(self,
in_channels,
encoder_channels,
num_filters,
relu_leakiness,
use_bn,
use_act=True,
kernel_size=3,
transposed_kernel_size=2,
num_layers=0,
mode='transposed',
padding='zero',
act_upsampling_only=False):
super(ConvDecodeUnit, self).__init__()
assert mode in ('transposed', 'nn', 'bilinear', 'pixelshuffle',
'nn-resize-conv', 'nn-biresize-conv')
use_bias = not use_bn or encoder_channels == 0
if mode == 'transposed':
upsample = [
nn.ConvTranspose2d(in_channels,
num_filters,
kernel_size=transposed_kernel_size,
stride=2,
bias=use_bias)
]
in_channels = num_filters
elif mode == 'nn':
upsample = [nn.Upsample(scale_factor=2, mode='nearest')]
elif mode == 'bilinear':
upsample = [nn.Upsample(scale_factor=2, mode='bilinear')]
elif mode == 'pixelshuffle':
upsample = [
get_same_padding_layer(kernel_size, stride=1, mode=padding),
nn.Conv2d(in_channels,
4 * num_filters,
kernel_size=kernel_size,
stride=1,
bias=use_bias),
nn.PixelShuffle(upscale_factor=2)
]
in_channels = num_filters
elif mode == 'nn-resize-conv' or mode == 'nn-biresize-conv':
resize_mode = 'nearest' if mode == 'nn-resize-conv' else 'bilinear'
upsample = [
nn.Upsample(scale_factor=2, mode=resize_mode),
get_same_padding_layer(kernel_size, stride=1, mode=padding),
nn.Conv2d(in_channels,
num_filters,
kernel_size=kernel_size,
stride=1,
bias=use_bias)
]
in_channels = num_filters
decode = []
if act_upsampling_only:
# Add batch norm and activation only on the upsampling path:
# This way saves some computation+memory by not activating on the
# encoded features again
if use_bn:
upsample += [nn.BatchNorm2d(in_channels)]
if use_act:
upsample += [nn.LeakyReLU(relu_leakiness, inplace=True)]
else:
# Add batch norm and activation for the concatenated features from both
# encoder path and upsampling path. Is kept here for legacy reasons.
if use_bn:
decode += [nn.BatchNorm2d(in_channels + encoder_channels)]
if use_act:
decode += [nn.LeakyReLU(relu_leakiness, inplace=True)]
if num_layers > 0:
decode += [
ConvEncodeUnit(in_channels + encoder_channels,
num_layers,
num_filters,
kernel_size,
relu_leakiness,
use_bn,
downsample=False,
use_act=use_act,
padding=padding)
]
self.upsample = nn.Sequential(*upsample)
self.decode = nn.Sequential(*decode)
def __init__(self, num_inputs, num_filters_per_layer, strides,
kernel_sizes=None, fc_layers=[], spatial_shape=None,
act_fn='lrelu', relu_leakiness=DEFAULT_RELU_LEAKINESS,
use_norm_layers=True, norm_layer='batch', use_weightnorm=False,
padding='zero', final_conv_kernel_size=1, use_biases=True,
final_average_pooling=False,
compute_features=False, dropout_after=[], dropout_prob=0.5):
"""Construct model
Parameters
----------
num_inputs: int
Number of input channels
num_filters_per_layer : list or tuple
Number of filters the discriminator uses in each layer
kernel_sizes : list
Shape of filters in each layer. Defaults to 3
strides : list
Strides of filters in each layer.
fc_layers : list
Number of channels of fully connected layers after convolutional layers.
If no fully connected layers are selected, the convolutional features
maps will be reduced to one dimension with a 1x1 convolution, and the
output is a probability map (corresponds to a PatchGAN)
spatial_shape : tuple
Spatial shape of input in the form of (height, width). Required if
using fully connected layers
act_fn : string
Activation function to use. Either `relu`, `prelu`, or `lrelu` (default)
relu_leakiness : float
If using lrelu, leakiness of the relus, if using prelu, initial value
for prelu parameters
use_norm_layers : bool or string
If true, use normalization layers. If `not-first`, skip the normalization
after the first convolutional layer
norm_layer : string
Normalization layer to use. `batch` for batch normalization or `instance`
for instance normalization
use_weightnorm : bool
If true, applies weight norm to all layers
padding : string
Type of padding to use. Either `zero`, `reflection`, or `replication`
final_conv_kernel_size : int
Shape of filter of final convolution, if using no fc layers.
Defaults to 1
final_average_pooling : bool
If true, reduce spatial dimensions to a scalar using an average pooling
layer. Only used if no fully connected layers were requested
use_biases : bool
If false, deactivate bias on all layers
compute_features : bool
If true, return the feature maps from forward pass in the key `features`
dropout_after : list
Indices of convolutional layers after which to insert layerwise dropout
dropout_prob : float
Probability to drop entries in dropout layers
"""
super(CNNDiscriminator, self).__init__()
if len(fc_layers) > 0:
assert spatial_shape is not None, \
'Need input spatial shape if using fully connected layers'
if kernel_sizes is None:
kernel_sizes = 3
if isinstance(kernel_sizes, int):
kernel_sizes = [kernel_sizes] * len(num_filters_per_layer)
assert len(num_filters_per_layer) == len(strides)
assert len(num_filters_per_layer) == len(kernel_sizes)
self.compute_features = compute_features
self.feature_layers = set()
in_channels = num_inputs
layer_idx = 0
layers = []
for num_filters, kernel_size, stride in zip(num_filters_per_layer,
kernel_sizes,
strides):
use_bias = use_biases and need_bias(use_norm_layers, norm_layer)
layers += (get_same_padding_layer(kernel_size=kernel_size, stride=stride,
mode=padding),
nn.Conv2d(in_channels, num_filters, kernel_size=kernel_size,
stride=stride,
bias=use_bias))
if use_norm_layers != 'not-first' and use_norm_layers:
layers.append(get_normalization_layer(norm_layer, num_filters))
elif use_norm_layers == 'not-first':
use_norm_layers = True
layers.append(get_activation_fn(act_fn, relu_leakiness, num_filters))
if compute_features:
self.feature_layers.add(layers[-1])
if layer_idx in dropout_after:
layers.append(nn.Dropout2d(p=dropout_prob, inplace=True))
in_channels = num_filters
layer_idx += 1
self.convs = nn.Sequential(*layers)
if len(fc_layers) > 0:
input_dims = self._infer_shape(self.convs, num_inputs, spatial_shape)
layers = []
for num_features in fc_layers[:-1]:
layers += (nn.Linear(input_dims, num_features, bias=use_biases),
get_activation_fn(act_fn, relu_leakiness, num_features))
input_dims = num_features
layers.append(nn.Linear(input_dims, fc_layers[-1]))
self.fcs = nn.Sequential(*layers)
self.final_conv = None
else:
self.fcs = None
final_conv = [nn.Conv2d(in_channels, out_channels=1,
kernel_size=final_conv_kernel_size, stride=1,
bias=use_biases)]
if final_average_pooling:
final_conv.append(nn.AdaptiveAvgPool2d((1, 1)))
self.final_conv = nn.Sequential(*final_conv)
def __init__(self, num_inputs, num_filters_per_layer, strides,
kernel_sizes=None, fc_layers=[], spatial_shape=None,
act_fn='lrelu', relu_leakiness=DEFAULT_RELU_LEAKINESS,
use_norm_layers=True, norm_layer='batch', padding='zero'):
"""Construct model
Parameters
----------
num_inputs: int
Number of input channels
num_filters_per_layer : list or tuple
Number of filters the discriminator uses in each layer
kernel_sizes : list
Shape of filters in each layer. Defaults to 3
strides : list
Strides of filters in each layer.
fc_layers : list
Number of channels of fully connected layers after convolutional layers.
If no fully connected layers are selected, the convolutional features
maps will be reduced to one dimension with a 1x1 convolution, and the
output is a probability map (corresponds to a PatchGAN)
spatial_shape : tuple
Spatial shape of input in the form of (height, width). Required if
using fully connected layers
act_fn : string
Activation function to use. Either `relu`, `prelu`, or `lrelu` (default)
relu_leakiness : float
If using lrelu, leakiness of the relus, if using prelu, initial value
for prelu parameters
use_norm_layers : bool or string
If true, use normalization layers. If `not-first`, skip the normalization
after the first convolutional layer
norm_layer : string
Normalization layer to use. `batch` for batch normalization or `instance`
for instance normalization
padding : string
Type of padding to use. Either `zero`, `reflection`, or `replication`
"""
super(CNNDiscriminator, self).__init__()
if len(fc_layers) > 0:
assert spatial_shape is not None, \
'Need input spatial shape if using fully connected layers'
if kernel_sizes is None:
kernel_sizes = 3
if isinstance(kernel_sizes, int):
kernel_sizes = [kernel_sizes] * len(num_filters_per_layer)
in_channels = num_inputs
model = []
for num_filters, kernel_size, stride in zip(num_filters_per_layer,
kernel_sizes,
strides):
model += (get_same_padding_layer(kernel_size=kernel_size, stride=stride,
mode=padding),
nn.Conv2d(in_channels, num_filters, kernel_size=kernel_size,
stride=stride,
bias=need_bias(use_norm_layers, norm_layer)))
if use_norm_layers != 'not-first' and use_norm_layers:
model.append(get_normalization_layer(norm_layer, num_filters))
elif use_norm_layers == 'not-first':
use_norm_layers = True
model.append(get_activation_fn(act_fn, relu_leakiness, num_filters))
in_channels = num_filters
self.convs = nn.Sequential(*model)
if len(fc_layers) > 0:
input_dims = self._infer_shape(self.convs, num_inputs, spatial_shape)
model = []
for num_features in fc_layers[:-1]:
model += (nn.Linear(input_dims, num_features),
get_activation_fn(act_fn, relu_leakiness, num_features))
input_dims = num_features
model.append(nn.Linear(input_dims, fc_layers[-1]))
self.fcs = nn.Sequential(*model)
self.final_conv = None
else:
self.fcs = None
self.final_conv = nn.Conv2d(in_channels, out_channels=1,
kernel_size=1, stride=1, bias=True)
def __init__(self, num_inputs, num_outputs, upscale_factor,
num_filters=64, num_res_blocks=16,
output_activation='tanh', act_fn='prelu',
relu_leakiness=DEFAULT_RELU_LEAKINESS,
use_norm_layers='not-first', norm_layer='batch',
padding='zero'):
"""Builds a SRResNet (Ledig et al, https://arxiv.org/abs/1609.04802)
Parameters
----------
num_inputs : int
Number of input channels
num_outputs : int
Number of output channels
upscale_factor : int
Factor by which the network upscales. Must be divisible by 2 or 3
num_filters : int
Number of convolutional filters to use
num_res_blocks : int
Number of residual blocks
output_activation : string
Either `softmax` or `tanh`. Activation function to use on the logits
act_fn : string
Activation function to use. Either `relu`, `prelu` (default), or `lrelu`
relu_leakiness : float
If using lrelu, leakiness of the relus, if using prelu, initial value
for prelu parameters
use_norm_layers : bool or string
If true, use normalization layers. If `not-first`, skip the normalization
after the first convolutional layer
norm_layer : string
Normalization layer to use. `batch` for batch normalization or `instance`
for instance normalization
padding : string
Type of padding to use. Either `zero`, `reflection`, or `replication`
"""
super(SRResNet, self).__init__()
upscale_factor = int(upscale_factor)
assert (upscale_factor == 1 or
upscale_factor % 2 == 0 or
upscale_factor % 3 == 0)
in_channels = num_inputs
initial_conv = [get_same_padding_layer(kernel_size=9, stride=1,
mode=padding),
nn.Conv2d(in_channels, num_filters, kernel_size=9,
stride=1,
bias=need_bias(use_norm_layers, norm_layer)),
get_activation_fn(act_fn, relu_leakiness, num_filters)]
in_channels = num_filters
if use_norm_layers != 'not-first' and use_norm_layers:
initial_conv.append(get_normalization_layer(norm_layer, in_channels))
elif use_norm_layers == 'not-first':
use_norm_layers = True
res_blocks = []
for idx in range(num_res_blocks):
res_blocks += [ResBlock(in_channels, num_filters, kernel_size=3,
use_norm_layers=use_norm_layers,
norm_layer=norm_layer, act_fn=act_fn,
relu_leakiness=relu_leakiness)]
second_conv = [get_same_padding_layer(kernel_size=3, stride=1,
mode=padding),
nn.Conv2d(in_channels, num_filters, kernel_size=3,
stride=1, bias=need_bias(use_norm_layers,
norm_layer))]
in_channels = num_filters
if use_norm_layers:
second_conv.append(get_normalization_layer(norm_layer, in_channels))
upsample = []
if upscale_factor > 1:
scale = 2 if upscale_factor % 2 == 0 else 3
for idx in range(upscale_factor // scale):
upsample += [get_same_padding_layer(kernel_size=3, stride=1,
mode=padding),
nn.Conv2d(in_channels, scale * scale * 256,
kernel_size=3, stride=1, bias=True),
nn.PixelShuffle(upscale_factor=scale),
get_activation_fn(act_fn, relu_leakiness, 256)]
in_channels = 256
final_conv = [get_same_padding_layer(kernel_size=9, stride=1,
mode=padding),
nn.Conv2d(in_channels, num_outputs, kernel_size=9,
stride=1, bias=True)]
if output_activation != 'none':
final_conv.append(get_activation_fn(output_activation))
self.initial_conv = nn.Sequential(*initial_conv)
self.body = nn.Sequential(*(res_blocks + second_conv))
self.upsample = nn.Sequential(*upsample)
self.final_conv = nn.Sequential(*final_conv)
self.output_activation = output_activation