def _build_model(self):
self.input_module = ConvBlock(
input_channels=self.input_channels,
output_channels=self.feature_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
)
if self.num_hidden_convolution_blocks > 0:
self.processing_module = ConvMultiBlock(
input_channels=self.feature_channels,
feature_channels=self.feature_channels,
output_channels=self.feature_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
batch_norm=self.batch_norm,
num_convolution_blocks=self.num_hidden_convolution_blocks
)
else:
self.processing_module = None
self.output_module = ConvBlock(
input_channels=self.feature_channels,
output_channels=self.output_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
)
def _build_model(self):
self.input_module = ConvBlock(
input_channels=self.input_channels,
output_channels=self.feature_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
)
if self.num_residual_blocks > 0:
self.processing_module = ResNetMultiBlock(
input_channels=self.feature_channels,
feature_channels=self.feature_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
num_resnet_blocks=self.num_residual_blocks,
num_hidden_layers=self.layers_per_residual_block,
)
else:
self.processing_module = None
self.output_module = ConvBlock(
input_channels=self.feature_channels,
output_channels=self.output_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
)
def _build_model(self):
if self.num_convolution_blocks_hr > 0:
self.convolutions_hr = ConvMultiBlock(
input_channels=self.input_channels,
feature_channels=self.feature_channels_hr,
output_channels=self.feature_channels_hr,
kernel_size=self.kernel_size_hr,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
batch_norm=self.batch_norm,
num_convolution_blocks=self.num_convolution_blocks_hr)
else:
self.convolutions_hr = None
self.downsampling_23 = ConvBlock(
input_channels=self.feature_channels_hr
if self.num_convolution_blocks_hr > 0 else self.input_channels,
output_channels=self.feature_channels_mr,
kernel_size=self.kernel_size_hr,
stride=(2, 3),
padding_mode=self.padding_mode,
use_batch_norm=self.batch_norm,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate)
if self.num_convolution_blocks_mr > 0:
self.convolutions_mr = ConvMultiBlock(
input_channels=self.feature_channels_mr,
feature_channels=self.feature_channels_mr,
output_channels=self.feature_channels_mr,
kernel_size=self.kernel_size_mr,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
batch_norm=self.batch_norm,
num_convolution_blocks=self.num_convolution_blocks)
else:
self.convolutions_mr = None
self.downsampling_21 = ConvBlock(
input_channels=self.feature_channels_mr,
output_channels=self.feature_channels_lr
if self.num_convolution_blocks_lr > 0 else self.output_channels,
kernel_size=self.kernel_size_mr,
stride=(2, 1),
padding_mode=self.padding_mode,
use_batch_norm=self.batch_norm,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate)
if self.num_convolution_blocks_lr > 0:
self.convolutions_lr = ConvMultiBlock(
input_channels=self.feature_channels_lr,
feature_channels=self.feature_channels_lr,
output_channels=self.output_channels,
kernel_size=self.kernel_size_lr,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
batch_norm=self.batch_norm,
num_convolution_blocks=self.num_convolution_blocks_lr)
else:
self.convolutions_lr = None
def _processing_module(self, input_channels, **kwargs):
if 'feature_channels' in kwargs.keys():
feature_channels = kwargs['feature_channels']
else:
feature_channels = input_channels
if self.layers_per_module == 0:
module = nn.Identity()
elif self.layers_per_module == 1:
module = ConvBlock(
input_channels=input_channels,
output_channels=input_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
use_batch_norm=self.batch_norm,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
)
else:
layers = [
ConvBlock(
input_channels=input_channels,
output_channels=feature_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
use_batch_norm=self.batch_norm,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
)
]
for i in range(self.layers_per_module - 2):
layers += [
ConvBlock(
input_channels=feature_channels,
output_channels=feature_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
use_batch_norm=self.batch_norm,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
)
]
layers += [
ConvBlock(
input_channels=feature_channels,
output_channels=input_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
use_batch_norm=self.batch_norm,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
)
]
module = nn.Sequential(*layers)
return module
def _build_init_module(self, input_channels, output_channels):
"""
Build an input module to process the data. For now, we assume that input has been processed
beforehand and we just deal with this processed input. No raw data processing is supported by default.
However, enabling init convolution enables the network to cope with raw data, as well.
Then, see that the input_channels match the number of input features.
"""
block = []
if self.init_upsampling:
# upsample image to final target resolution
block += [
nn.Upsample(scale_factor=(4, 3), mode=self.interpolation_mode)
]
if self.init_convolution:
# process the input with a convolution block
block += ConvBlock(input_channels=input_channels,
output_channels=output_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm)
self.upsample_input = nn.Sequential(*block)
def _upsampling_feature_module(self, input_channels, output_channels,
scale_factor, **kwargs):
module = ConvBlock(input_channels=input_channels,
output_channels=output_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm)
return module
def _build_model(self):
self._build_hr_module()
self.downsampling_23 = ConvBlock(
input_channels=self.feature_channels_hr,
output_channels=self.feature_channels_mr,
kernel_size=self.kernel_size_hr,
stride=(2, 3),
padding_mode=self.padding_mode,
use_batch_norm=self.batch_norm,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate)
self._build_mr_module()
self.downsampling_21 = ConvBlock(
input_channels=self.feature_channels_mr,
output_channels=self.feature_channels_lr,
kernel_size=self.kernel_size_mr,
stride=(2, 1),
padding_mode=self.padding_mode,
use_batch_norm=self.batch_norm,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate)
self._build_lr_module()
def _build_model(self):
self.upsampling_module = nn.Sequential(
ResamplingBlock2D(input_channels=self.input_channels,
scale_factor=self.scale_factor,
interpolation_mode=self.interpolation_mode),
ConvBlock(
input_channels=self.input_channels,
output_channels=self.feature_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
),
)
if self.num_convolution_blocks > 0:
self.processing_module = ConvMultiBlock(
input_channels=self.feature_channels,
feature_channels=self.feature_channels,
output_channels=self.feature_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
batch_norm=self.batch_norm,
num_convolution_blocks=self.num_convolution_blocks)
else:
self.processing_module = None
self.output_module = ConvBlock(
input_channels=self.feature_channels,
output_channels=self.output_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope if self.output_activation else 1.0,
dropout_rate=self.dropout_rate if self.output_activation else 0.0,
use_batch_norm=self.batch_norm
if self.output_activation else False,
)
def _build_lr_module(self):
if self.num_inception_blocks_lr > 0:
layers = []
for _ in range(self.num_inception_blocks_lr):
layers += [
InceptionBlock(
input_channels=self.feature_channels_lr,
feature_channels=self.feature_channels_lr,
output_channels=self.feature_channels_lr,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
)
]
layers += [
ConvBlock(
input_channels=self.feature_channels_lr,
output_channels=self.output_channels,
kernel_size=self.kernel_size_lr,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
)
]
self.module_lr = nn.Sequential(*layers)
else:
self.module_lr = ConvBlock(
input_channels=self.feature_channels_lr,
output_channels=self.output_channels,
kernel_size=self.kernel_size_lr,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
)
def _build_model(self):
self.input_module = ConvBlock(
input_channels=self.input_channels,
output_channels=self.feature_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
)
if self.num_residual_blocks > 0:
layers = []
for _ in range(self.num_inception_blocks):
layers += [
InceptionBlock(
input_channels=self.feature_channels,
feature_channels=self.feature_channels,
output_channels=self.feature_channels,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
)
]
self.processing_module = nn.Sequential(*layers)
else:
self.processing_module = None
self.output_module = ConvBlock(
input_channels=self.feature_channels,
output_channels=self.output_channels,
kernel_size=self.kernel_size,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
)
def _build_model(self):
if self.num_residual_blocks_hr > 0:
self.module_hr = nn.Sequential(
ConvBlock(
input_channels=self.input_channels,
output_channels=self.feature_channels_hr,
kernel_size=self.kernel_size_hr,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
),
ResNetMultiBlock(
input_channels=self.feature_channels_hr,
feature_channels=self.feature_channels_hr,
kernel_size=self.kernel_size_hr,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
num_resnet_blocks=self.num_residual_blocks_hr,
num_hidden_layers=self.layers_per_residual_block,
)
)
else:
self.module_hr = ConvBlock(
input_channels=self.input_channels,
output_channels=self.feature_channels_hr,
kernel_size=self.kernel_size_hr,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
)
self.downsampling_23 = ConvBlock(
input_channels=self.feature_channels_hr,
output_channels=self.feature_channels_mr,
kernel_size=self.kernel_size_hr,
stride=(2, 3),
padding_mode=self.padding_mode,
use_batch_norm=self.batch_norm,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate
)
if self.num_residual_blocks_mr > 0:
self.module_mr = ResNetMultiBlock(
input_channels=self.feature_channels_mr,
feature_channels=self.feature_channels_mr,
kernel_size=self.kernel_size_mr,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
num_resnet_blocks=self.num_residual_blocks_mr,
num_hidden_layers=self.layers_per_residual_block,
)
else:
self.module_mr = None
self.downsampling_21 = ConvBlock(
input_channels=self.feature_channels_mr,
output_channels=self.feature_channels_lr,
kernel_size=self.kernel_size_mr,
stride=(2, 1),
padding_mode=self.padding_mode,
use_batch_norm=self.batch_norm,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate
)
if self.num_residual_blocks_lr > 0:
self.module_lr = nn.Sequential(
ResNetMultiBlock(
input_channels=self.feature_channels_lr,
feature_channels=self.feature_channels_lr,
kernel_size=self.kernel_size_lr,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
num_resnet_blocks=self.num_residual_blocks_lr,
num_hidden_layers = self.layers_per_residual_block,
),
ConvBlock(
input_channels=self.feature_channels_lr,
output_channels=self.output_channels,
kernel_size=self.kernel_size_lr,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
)
)
else:
self.module_lr = ConvBlock(
input_channels=self.feature_channels_lr,
output_channels=self.output_channels,
kernel_size=self.kernel_size_lr,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
)
def __init__(self, **kwargs):
super(self.__class__, self).__init__(**kwargs)
if self.feature_channels is None:
self.feature_channels = 64
self._require_not_none('input_channels', 'output_channels')
self.model = nn.Sequential(
ResNetMultiBlock(input_channels=self.input_channels,
feature_channels=self.input_channels,
output_channels=self.input_channels,
kernel_size=3,
padding_mode=self.padding_mode,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate,
use_batch_norm=self.batch_norm,
num_resnet_blocks=self.num_residual_blocks,
num_hidden_layers=self.layers_per_residual_block -
1,
output_activation=True),
ResamplingBlock2D(input_channels=self.input_channels,
output_channels=self.input_channels,
size=None,
scale_factor=(2, 1),
interpolation_mode='nearest'),
ConvBlock(input_channels=self.input_channels,
output_channels=self.input_channels,
kernel_size=3,
padding_mode='replication',
use_batch_norm=True,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate),
ResamplingBlock2D(input_channels=self.input_channels,
output_channels=self.input_channels,
size=None,
scale_factor=(1, 3),
interpolation_mode='nearest'),
ConvBlock(
input_channels=self.input_channels,
output_channels=self.input_channels,
# kernel_size=[3, 5],
kernel_size=[3, 5],
padding_mode='replication',
use_batch_norm=True,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate),
ResamplingBlock2D(input_channels=self.input_channels,
output_channels=self.input_channels,
size=None,
scale_factor=(2, 1),
interpolation_mode='nearest'),
#####################################################################
# Added by michael --> this should be exactly like the paper 1x Conv, ReLU after supersampling
ConvBlock(input_channels=self.input_channels,
output_channels=self.input_channels,
kernel_size=3,
padding_mode='replication',
use_batch_norm=True,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate),
# + 1x Conv, ReLU afterwards
ConvBlock(input_channels=self.input_channels,
output_channels=self.input_channels,
kernel_size=3,
padding_mode='replication',
use_batch_norm=True,
leaky_slope=self.leaky_slope,
dropout_rate=self.dropout_rate),
#####################################################################
ConvBlock(input_channels=self.input_channels,
output_channels=self.output_channels,
kernel_size=3,
padding_mode='replication',
use_batch_norm=False,
leaky_slope=1.0,
dropout_rate=0.0),
)
