def __init__(self, opt):
super(Generator, self).__init__()
# Read options
nonlinear_layer = get_nonlinear_layer(opt.nonlinearity)
norm_layer = get_norm_layer(opt.norm)
depth = int(log(opt.image_size // opt.latent_size, 2))
layers = []
# Deconde input into latent
in_channels = len(opt.input_idx.split(','))
assert opt.in_channels > in_channels, 'in_channels must be > num input_idx'
out_channels = min(opt.num_channels * 2**depth, opt.max_channels)
bias = norm_layer == Identity or norm_layer == nn.InstanceNorm2d
layers += [
nn.Linear(opt.in_channels,
out_channels * opt.latent_size**2,
bias=bias),
View(opt.latent_size),
norm_layer(out_channels),
nonlinear_layer(True)
]
# Upsampling decoding blocks
for i in range(depth):
in_channels = out_channels
out_channels = min(opt.num_channels * 2**(depth - i - 1),
opt.max_channels)
layers += get_conv_block(in_channels, out_channels,
nonlinear_layer, norm_layer, 'up', False)
in_channels = out_channels
layers += [nn.Conv2d(in_channels, 1, 3, 1, 1)]
self.block = nn.Sequential(*layers)
self.apply(weights_init)
def __init__(self, opt):
super(Discriminator, self).__init__()
# Read options
nonlinear_layer = get_nonlinear_layer('LeakyReLU')
norm = opt.norm if opt.norm_dis else 'None'
norm_layer = get_norm_layer(norm)
depth = int(log(opt.image_size // opt.latent_size, 2))
assert depth >= 1, 'image_size must be >= 8'
in_channels = 1
out_channels = opt.num_channels
layers = get_conv_block(in_channels, out_channels, nonlinear_layer,
norm_layer, 'none', False, opt.kernel_size)
# Get all model blocks
for i in range(depth):
# Set number of channels for conv
in_channels = out_channels
out_channels = min(out_channels * 2, opt.max_channels)
# Define downsampling block
layers += get_conv_block(in_channels, out_channels,
nonlinear_layer, norm_layer, 'down',
False, opt.kernel_size)
# Output single number pred
layers += [
View(),
nn.Linear(out_channels * opt.latent_size**2, opt.num_preds)
]
self.block = nn.Sequential(*layers)
# Initialize weights
self.apply(weights_init)
def __init__(self, opt, type_name):
"""Initialize discriminator network.
This method creates and initializes trainable discriminator layers.
Args:
opt: Options specified by the user.
type_name: Specified discriminator type name.
"""
super(Discriminator, self).__init__()
# Read options
self.input_sizes = opt['%s_input_sizes' % type_name].split(',')
self.input_sizes = [int(i) for i in self.input_sizes]
output_sizes = opt['%s_output_sizes' % type_name].split(',')
self.output_sizes = [int(i) for i in output_sizes]
output_weights = opt['%s_output_weights' % type_name].split(',')
self.output_weights = [float(i) for i in output_weights]
kernel_size = opt['%s_kernel_size' % type_name]
kernel_size_io = opt['%s_kernel_size_io' % type_name]
input_num_channels = opt['%s_input_num_channels' %
type_name].split(',')
input_num_channels = [int(i) for i in input_num_channels]
num_channels = opt['%s_num_channels' % type_name]
max_channels = opt['%s_max_channels' % type_name]
self.adv_loss_type = opt['%s_adv_loss_type' % type_name]
self.use_encoder = opt['%s_use_encoder' % type_name]
norm_layer = utils.get_norm_layer(opt['%s_norm_layer' % type_name])
norm_layer_cat = utils.get_norm_layer(opt['%s_norm_layer_cat' %
type_name])
nonlinear_layer = utils.get_nonlinear_layer(opt['%s_nonlinear_layer' %
type_name])
aux_channels = opt['aux_channels']
# Setup number of channels
if not num_channels: num_channels = opt['num_channels']
if not max_channels: max_channels = opt['max_channels']
# Calculate final tensor spatial size
if (self.output_sizes[-1] == 2
or len(self.output_sizes) > 1 and self.output_sizes[-2] < 4):
result_spatial_size = self.output_sizes[-2]
else:
result_spatial_size = 4
# Calculate network depth
self.depth = int(log(self.input_sizes[0] // result_spatial_size, 2))
# Calculate output shape
self.output_shapes = []
for s in self.output_sizes:
self.output_shapes += [torch.Size([opt['batch_size'], s**2])]
# Initialize network
self.blocks = nn.ModuleList()
self.concat_blocks = nn.ModuleList()
self.output_blocks = nn.ModuleList()
in_channels = input_num_channels[0]
out_channels = num_channels
current_size = self.input_sizes[0]
self.concat_blocks_depth = []
self.output_blocks_depth = []
for i in range(self.depth):
# Define downsampling block
self.blocks += [
utils.get_conv_block(in_channels=in_channels,
out_channels=out_channels,
nonlinear_layer=nonlinear_layer,
norm_layer=norm_layer
if i else utils.get_norm_layer('none'),
mode='down',
sequential=True,
kernel_size=kernel_size)
]
current_size //= 2
in_channels = out_channels
if current_size in self.input_sizes:
k = self.input_sizes.index(current_size)
# Define concat block
self.concat_blocks += [
utils.ConcatBlock(enc_channels=input_num_channels[k],
out_channels=in_channels,
nonlinear_layer=nonlinear_layer,
norm_layer=norm_layer,
norm_layer_cat=norm_layer_cat,
kernel_size=kernel_size_io)
]
self.concat_blocks_depth += [i]
in_channels *= 2
if current_size in self.output_sizes:
# Define PatchGAN output block
self.output_blocks += [
nn.Sequential(
nn.Conv2d(in_channels=in_channels,
out_channels=1,
kernel_size=kernel_size_io,
stride=1,
padding=kernel_size_io // 2,
bias=False), utils.View())
]
self.output_blocks_depth += [i]
out_channels = min(out_channels * 2, max_channels)
if 1 in self.output_sizes:
# Final probability prediction
self.output_blocks += [
nn.Sequential(
nn.Conv2d(in_channels=out_channels,
out_channels=1,
kernel_size=current_size,
stride=current_size,
padding=0,
bias=False), utils.View())
]
# Auxiliary prediction
#if aux_channels and opt['dis_aux_loss_weight']:
# self.aux_output_block = nn.Sequential(nn.Conv2d(
# in_channels=out_channels,
# out_channels=aux_channels,
# kernel_size=current_size,
# stride=current_size,
# padding=0,
# bias=False),
# utils.View())
# Initialize weights
self.apply(utils.weights_init)
def __init__(self, opt):
super(Generator, self).__init__()
# Get constructors of the required modules
norm_layer = utils.get_norm_layer(opt.gen_norm_layer)
upsampling_layer = utils.get_upsampling_layer(opt.gen_upsampling_layer)
# Calculate the amount of downsampling and upsampling convolutional blocks
num_down_blocks = int(log(opt.image_size // opt.gen_latent_size, 2))
num_up_blocks = int(log(opt.image_size // opt.gen_latent_size, 2))
# Read parameters for convolutional blocks
in_channels = opt.gen_num_channels
padding = (opt.gen_kernel_size - 1) // 2
bias = norm_layer != nn.BatchNorm2d
# First block is without normalization
layers = [
nn.Conv2d(3, in_channels, 7, 1, 3, bias=False),
nn.ReLU(True)
]
# Downsampling blocks
for i in range(num_down_blocks):
# Increase the number of channels by 2x
out_channels = min(in_channels * 2, opt.gen_max_channels)
layers += [
nn.Conv2d(in_channels, out_channels, opt.gen_kernel_size, 2,
padding, bias),
norm_layer(out_channels),
nn.ReLU(True)
]
in_channels = out_channels
# Residual blocks
for i in range(opt.gen_num_res_blocks):
layers += [utils.ResBlock(in_channels, norm_layer)]
# Upsampling blocks
for i in range(num_up_blocks):
# Decrease the number of channels by 2x
out_channels = opt.gen_num_channels * 2**(num_up_blocks - i - 1)
out_channels = max(min(out_channels, opt.gen_max_channels),
opt.gen_num_channels)
layers += upsampling_layer(in_channels, out_channels,
opt.gen_kernel_size, 2, bias)
layers += [norm_layer(out_channels), nn.ReLU(True)]
in_channels = out_channels
# Last block outputs values in range [-1, 1]
layers += [nn.Conv2d(out_channels, 3, 7, 1, 3, bias=False), nn.Tanh()]
self.generator = nn.Sequential(*layers)
# Initialize weights
self.apply(utils.weights_init)
def __init__(self, opt, domain, type_name):
super(Generator, self).__init__()
opt = vars(opt)
# Check which mapping does this generator perform
ZtoB = domain == 'B'
# Read options
output_size = opt['img_size_%s' % 'B' if ZtoB else 'A']
output_channels = opt['img_channels_%s' % 'B' if ZtoB else 'A']
self.noise_channels = opt['%s_noise_channels' % type_name]
num_channels = opt['%s_num_channels' % type_name]
max_channels = opt['%s_max_channels' % type_name]
num_layers = opt['%s_num_layers' % type_name]
kernel_size = opt['%s_kernel_size' % type_name]
norm_layer = utils.get_norm_layer(opt['%s_norm_layer' % type_name])
norm_layer_lin = utils.get_norm_layer(opt['%s_norm_layer' % type_name],
dims=1)
norm_layer_cat = utils.get_norm_layer(opt['%s_norm_layer_cat' %
type_name])
upsampling_layer = opt['%s_upsampling_layer' % type_name]
nonlinear_layer = utils.get_nonlinear_layer(opt['%s_nonlinear_layer' %
type_name])
final_nonlinear_layer = utils.get_nonlinear_layer(
opt['%s_output_range' % type_name])
aux_channels = opt['aux_channels']
# Setup number of channels
if not num_channels: num_channels = opt['num_channels']
if not max_channels: max_channels = opt['max_channels']
# Calculate network depth
depth = int(log(output_size // 4, 2))
out_channels = min(num_channels * 2**depth, max_channels)
# Encoder aux
if aux_channels:
out_channels //= 2
self.encoder_aux = utils.get_linear_block(
in_channels=aux_channels,
out_channels=out_channels,
nonlinear_layer=nonlinear_layer,
norm_layer=norm_layer_lin,
sequential=True)
# Encoder noise
self.encoder_noise = utils.get_linear_block(
in_channels=self.noise_channels,
out_channels=out_channels,
nonlinear_layer=nonlinear_layer,
norm_layer=norm_layer_lin,
sequential=True)
if aux_channels:
out_channels *= 2
in_channels = out_channels
# Build upsampling blocks
layers = [utils.View(), norm_layer_cat(out_channels)]
layers += utils.get_upsampling_block(in_channels=in_channels,
out_channels=out_channels,
nonlinear_layer=nonlinear_layer,
norm_layer=norm_layer,
mode=upsampling_layer,
sequential=False,
kernel_size=kernel_size,
num_layers=num_layers,
factor=4)
for i in range(depth):
in_channels = out_channels
out_channels = min(num_channels * 2**(depth - 1 - i), max_channels)
layers += utils.get_upsampling_block(
in_channels=in_channels,
out_channels=out_channels,
nonlinear_layer=nonlinear_layer,
norm_layer=norm_layer,
mode=upsampling_layer,
sequential=False,
kernel_size=kernel_size,
num_layers=num_layers)
layers += [
nn.Conv2d(in_channels=out_channels,
out_channels=output_channels,
kernel_size=7,
stride=1,
padding=3,
bias=False),
final_nonlinear_layer(True)
]
self.generator = nn.Sequential(*layers)
# Initialize weights
self.apply(utils.weights_init)