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)