def forward(self, c_img, s_img, a, batch_size): c_img_temp = self.vgg_norm(c_img, batch_size) s_img_temp = self.vgg_norm(s_img, batch_size) c_feature = self.encoder_1(c_img_temp) c_feature = self.encoder_2(c_feature) c_feature = self.encoder_3(c_feature) c_feature = self.encoder_4(c_feature) s_feature_1 = self.encoder_1(s_img_temp) s_feature_2 = self.encoder_2(s_feature_1) s_feature_3 = self.encoder_3(s_feature_2) s_feature_4 = self.encoder_4(s_feature_3) a_feature = self.adain(c_feature, s_feature_4) t = (1-a) * c_feature + a * a_feature result = self.decoder(t) result_temp = self.vgg_norm(result, batch_size) r_feature_1 = self.encoder_1(result_temp) r_feature_2 = self.encoder_2(r_feature_1) r_feature_3 = self.encoder_3(r_feature_2) r_feature_4 = self.encoder_4(r_feature_3) s_loss_1 = StyleLoss(s_feature_1, r_feature_1) s_loss_2 = StyleLoss(s_feature_2, r_feature_2) s_loss_3 = StyleLoss(s_feature_3, r_feature_3) s_loss_4 = StyleLoss(s_feature_4, r_feature_4) s_loss = s_loss_1 + s_loss_2 + s_loss_3 + s_loss_4 return a_feature, r_feature_4, result, s_loss
def __init__(self, config): super(InpaintModel, self).__init__('InpaintModel', config) #generator = Generator() generator = Generator_SE() discriminator = Discriminator(in_channels=3, use_sigmoid=False) # data = torch.load('ablation_v0/InpaintModel121_gen.pth') # generator.load_state_dict(data['generator']) # self.iteration = data['iteration'] # data = torch.load('ablation_v0/InpaintModel121_dis.pth') # discriminator.load_state_dict(data['discriminator']) l1_loss = nn.L1Loss() adversarial_loss = AdversarialLoss(type='hinge') perceptual_loss = PerceptualLoss() style_loss = StyleLoss() self.add_module('generator', generator) self.add_module('discriminator', discriminator) self.add_module('l1_loss', l1_loss) self.add_module('adversarial_loss', adversarial_loss) self.add_module('perceptual_loss', perceptual_loss) self.add_module('style_loss', style_loss) self.optimizer = optim.Adam(params=generator.parameters(), lr=config.LR) self.dis_optimizer = optim.Adam(params=discriminator.parameters(), lr=config.LR * config.D2G_LR)
def get_style_model_and_losses(cnn, device, normalization_mean, normalization_std, style_img, content_img, content_layers=content_layers_default, style_layers=style_layers_default): cnn = copy.deepcopy(cnn) # normalization module normalization = Normalization(normalization_mean, normalization_std).to(device) # just in order to have an iterable access to or list of content/syle # losses content_losses = [] style_losses = [] # assuming that cnn is a nn.Sequential, so we make a new nn.Sequential # to put in modules that are supposed to be activated sequentially model = nn.Sequential(normalization) i = 0 # increment every time we see a conv for layer in cnn.children(): if isinstance(layer, nn.Conv2d): i += 1 name = 'conv_{}'.format(i) elif isinstance(layer, nn.ReLU): name = 'relu_{}'.format(i) # The in-place version doesn't play very nicely with the ContentLoss # and StyleLoss we insert below. So we replace with out-of-place # ones here. layer = nn.ReLU(inplace=False) elif isinstance(layer, nn.MaxPool2d): name = 'pool_{}'.format(i) elif isinstance(layer, nn.BatchNorm2d): name = 'bn_{}'.format(i) else: raise RuntimeError('Unrecognized layer: {}'.format(layer.__class__.__name__)) model.add_module(name, layer) if name in content_layers: # add content loss: target = model(content_img).detach() content_loss = ContentLoss(target) model.add_module("content_loss_{}".format(i), content_loss) content_losses.append(content_loss) if name in style_layers: # add style loss: target_feature = model(style_img).detach() style_loss = StyleLoss(target_feature) model.add_module("style_loss_{}".format(i), style_loss) style_losses.append(style_loss) # now we trim off the layers after the last content and style losses for i in range(len(model) - 1, -1, -1): if isinstance(model[i], ContentLoss) or isinstance(model[i], StyleLoss): break model = model[:(i + 1)] return model, style_losses, content_losses
def __init__(self, config): super(InpaintingModel, self).__init__() self.name = 'InpaintingModel' self.config = config self.iteration = 0 self.gen_weights_path = os.path.join(config.save_model_dir, 'InpaintingModel_gen.pth') self.dis_weights_path = os.path.join(config.save_model_dir, 'InpaintingModel_dis.pth') self.generator = FRRNet() self.discriminator = Discriminator(in_channels=3, use_sigmoid=True) if torch.cuda.device_count() > 1: device_ids=range(torch.cuda.device_count()) self.generator = nn.DataParallel(self.generator, device_ids) self.discriminator = nn.DataParallel(self.discriminator , device_ids) self.l1_loss = nn.L1Loss() self.l2_loss = nn.MSELoss() self.style_loss = StyleLoss() self.adversarial_loss = AdversarialLoss() self.gen_optimizer = optim.Adam( params=self.generator.parameters(), lr=float(config.LR), betas=(0.0, 0.9) ) self.dis_optimizer = optim.Adam( params=self.discriminator.parameters(), lr=float(config.LR) * float(config.D2G_LR), betas=(0.0, 0.9) )
def _attempt_insert_style_loss(self): if self.layer_count in self.style_layer: self.style_layer.remove(self.layer_count) target_features = self._model(self.style_tensor).detach() layer = StyleLoss(target_features) self._model.add_module('style_loss_{}'.format(self.layer_count), layer) self._style_losses.append(layer)
def get_style_model_and_losses( cnn, cnn_normalization_mean, cnn_normalization_std, style_img, content_img, ): cnn = copy.deepcopy(cnn) content_layers = ['conv_4'] style_layers = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5'] normalization = Normalization(cnn_normalization_mean, cnn_normalization_std).to(device) content_losses = [] style_losses = [] model = nn.Sequential(normalization) i = 0 for layer in cnn.children(): if isinstance(layer, nn.Conv2d): i += 1 name = 'conv_{}'.format(i) elif isinstance(layer, nn.ReLU): name = 'relu_{}'.format(i) layer = nn.ReLU(inplace=False) elif isinstance(layer, nn.MaxPool2d): name = 'pool_{}'.format(i) elif isinstance(layer, nn.BatchNorm2d): name = 'bn_{}'.format(i) else: raise RuntimeError('Unrecognized layer: {}'.format( layer.__class__.__name__)) model.add_module(name, layer) if name in content_layers: target = model(content_img).detach() content_loss = ContentLoss(target) model.add_module('content_loss_{}'.format(i), content_loss) content_losses.append(content_loss) if name in style_layers: target_feature = model(style_img).detach() style_loss = StyleLoss(target_feature) model.add_module('style_loss_{}'.format(i), style_loss) style_losses.append(style_loss) for i in range(len(model) - 1, -1, -1): if isinstance(model[i], ContentLoss) or isinstance( model[i], StyleLoss): break model = model[:(i + 1)] return model, style_losses, content_losses
def get_style_model_and_losses(parent_model, style_img, content_img, content_layers=content_layers_default, style_layers=style_layers_default): parent_model = copy.deepcopy(parent_model) content_losses = [] style_losses = [] model = nn.Sequential(normalization) i = 0 # increment every time we see a conv for layer in parent_model.children(): if isinstance(layer, nn.Conv2d): i += 1 name = 'conv_{}'.format(i) elif isinstance(layer, nn.ReLU): name = 'relu_{}'.format(i) layer = nn.ReLU(inplace=False) elif isinstance(layer, nn.MaxPool2d): name = 'pool_{}'.format(i) elif isinstance(layer, nn.BatchNorm2d): name = 'bn_{}'.format(i) else: raise RuntimeError('Unrecognized layer: {}' .format(layer.__class__.__name__)) model.add_module(name, layer) if name in content_layers: # add content loss: target = model(content_img).detach() content_loss = ContentLoss(target) model.add_module("content_loss_{}".format(i), content_loss) content_losses.append(content_loss) if name in style_layers: # add style loss: target_feature = model(style_img).detach() style_loss = StyleLoss(target_feature) model.add_module("style_loss_{}".format(i), style_loss) style_losses.append(style_loss) for i in range(len(model) - 1, -1, -1): if isinstance(model[i], ContentLoss) or \ isinstance(model[i], StyleLoss): break model = model[:(i + 1)] return model, style_losses, content_losses
def build_model(pretrained, style, content, device, content_layers=[4], style_layers=[0, 1, 2, 3, 4, 5, 6, 7]): '''build style transfer model from pretrained model ''' model = nn.Sequential() count = 0 # use device style = style.to(device) content = content.to(device) pretrained = pretrained.to(device) # max cnn layer max_layer = max(max(content_layers), max(style_layers)) # gets weird color pixels if not noralized model.add_module('normalize', Normalize(device)) # loop through pretrained model and add loss for idx, layer in enumerate(pretrained.children()): if isinstance(layer, nn.Conv2d): # return model if max layer reached if count > max_layer: continue model.add_module(f"Conv2D_{count}", layer) count += 1 # add content loss if count in content_layers: content_output = model(content).detach() model.add_module(f"Content_loss", ContentLoss(content_output)) # add style loss if count in style_layers: style_output = model(style).detach() model.add_module(f"Loss_{count}", StyleLoss(style_output)) else: model.add_module(f"{idx}", layer) return model
def __init__(self, config): super().__init__('SRModel', config) # generator input: [gray(1) + edge(1)] # discriminator input: [gray(1)] generator = SRGenerator() discriminator = Discriminator( in_channels=1, use_sigmoid=config.GAN_LOSS != 'hinge') # 3-->1 if len(config.GPU) > 1: generator = nn.DataParallel(generator, config.GPU) discriminator = nn.DataParallel(discriminator, config.GPU) l1_loss = nn.L1Loss() content_loss = ContentLoss() style_loss = StyleLoss() adversarial_loss = AdversarialLoss(type=config.GAN_LOSS) kernel = np.zeros((self.config.SCALE, self.config.SCALE)) kernel[0, 0] = 1 #kernel_weight = torch.tensor(np.tile(kernel, (3, 1, 1, 1))).float().to(config.DEVICE) # (out_channels, in_channels/groups, height, width) #self.add_module('scale_kernel', kernel_weight) #self.scale_kernel = torch.tensor(np.tile(kernel, (1, 1, 1, 1))).float().to(config.DEVICE) #3-->1 self.add_module('generator', generator) self.add_module('discriminator', discriminator) self.add_module('l1_loss', l1_loss) self.add_module('content_loss', content_loss) self.add_module('style_loss', style_loss) self.add_module('adversarial_loss', adversarial_loss) self.gen_optimizer = optim.Adam(params=generator.parameters(), lr=float(config.LR), betas=(config.BETA1, config.BETA2)) self.dis_optimizer = optim.Adam(params=discriminator.parameters(), lr=float(config.LR), betas=(config.BETA1, config.BETA2))
def get_style_model_and_loss(base_model, base_mean, base_std, style_img, content_img, content_layers=content_layers, style_layers=style_layers): base_model = copy.deepcopy( base_model) # Separate Object that doesn't affect each other norm = Normalization(base_mean, base_std).to(device) content_losses, style_losses = [], [] model = nn.Sequential(norm) # First Layer = Normalization Layer i = 0 # Count CNN layers for layer in base_model.children(): if isinstance(layer, nn.Conv2d): i += 1 name = "conv_{}".format(i) elif isinstance(layer, nn.ReLU): name = "relu_{}".format(i) layer = nn.ReLU(inplace=False) elif isinstance(layer, nn.MaxPool2d): name = "pool_{}".format(i) elif isinstance(layer, nn.BatchNorm2d): name = "bn_{}".format(i) else: raise RuntimeError("Unrecognized layer: {}".format( layer.__class__.__name__)) model.add_module( name, layer ) # Sequentially stack layers of VGG to our new model (Copy most of them, and insert Losses in the right place) if name in content_layers: target = model( content_img).detach() # Feature map of content img so far content_loss = ContentLoss(target) # input is directly fed model.add_module( "content_loss_{}".format(i), content_loss ) # Add a layer that computes loss and returns the original input (Like identity operation in a sense) content_losses.append(content_loss) if name in style_layers: target_feature = model(style_img).detach() style_loss = StyleLoss(target_feature) model.add_module( "style_loss_{}".format(i), style_loss ) # Again, compute the gradient and returns input as is style_losses.append(style_loss) # Get rid of unnecessary layers after style and content loss for i in range(len(model) - 1, -1, -1): if isinstance(model[i], ContentLoss) or isinstance( model[i], StyleLoss): break model = model[:(i + 1)] return model, style_losses, content_losses
def get_style_model_and_losses(cnn, style_img, content_img, style_weight=1000, content_weight=1, content_layers=content_layers_default, style_layers=style_layers_default): cnn = copy.deepcopy(cnn) # just in order to have an iterable access to or list of content/syle # losses content_losses = [] style_losses = [] model = nn.Sequential() # the new Sequential module network gram = GramMatrix( ) # we need a gram module in order to compute style targets # move these modules to the GPU if possible: if use_cuda: model = model.cuda() gram = gram.cuda() i = 1 for layer in list(cnn): if isinstance(layer, nn.Conv2d): name = "conv_" + str(i) model.add_module(name, layer) if name in content_layers: # add content loss: target = model(content_img).clone() content_loss = ContentLoss(target, content_weight) model.add_module("content_loss_" + str(i), content_loss) content_losses.append(content_loss) if name in style_layers: # add style loss: target_feature = model(style_img).clone() target_feature_gram = gram(target_feature) style_loss = StyleLoss(target_feature_gram, style_weight) model.add_module("style_loss_" + str(i), style_loss) style_losses.append(style_loss) if isinstance(layer, nn.ReLU): name = "relu_" + str(i) model.add_module(name, layer) if name in content_layers: # add content loss: target = model(content_img).clone() content_loss = ContentLoss(target, content_weight) model.add_module("content_loss_" + str(i), content_loss) content_losses.append(content_loss) if name in style_layers: # add style loss: target_feature = model(style_img).clone() target_feature_gram = gram(target_feature) style_loss = StyleLoss(target_feature_gram, style_weight) model.add_module("style_loss_" + str(i), style_loss) style_losses.append(style_loss) i += 1 if isinstance(layer, nn.MaxPool2d): name = "pool_" + str(i) model.add_module(name, layer) # *** return model, style_losses, content_losses
def get_style_model_and_losses( model, style_img, content_img, device, normalization_mean=[0.485, 0.456, 0.406], normalization_std=[0.229, 0.224, 0.225], content_layers=['conv_4'], style_layers=['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']): """Create a new model from pretrained model by adding content loss and style loss layers. Parameters ---------- model(torchvision model): pretrained model. In NST paper VGG19 is used. style_img (tensor): style image content_img (tensor): content image device (str): device to run model normalization_mean (list): default mean of VGG networks normalization_std (list): default standard deviation of VGG networks content_layers (list): add content loss after the convolutional layers are detected style_layers (list): add style loss after the convolutional layers are detected """ cnn = model.features.to(device).eval() cnn = copy.deepcopy( cnn ) # for more information, refer https://www.programiz.com/python-programming/shallow-deep-copy # normalization module normalization_mean = torch.tensor(normalization_mean).to(device) normalization_std = torch.tensor(normalization_std).to(device) normalization = Normalization(normalization_mean, normalization_std).to(device) content_losses = [] style_losses = [] # assuming that cnn is a nn.Sequential, so we make a new nn.sequential to put in # modules that are supposed to be activated sequentially model = nn.Sequential(normalization) i = 0 # increment every time we see a conv for layer in cnn.children(): if isinstance(layer, nn.Conv2d): i += 1 name = 'conv_{}'.format(i) elif isinstance(layer, nn.ReLU): name = 'relu_{}'.format(i) # The in-place version doesn't play very nicely with the ContentLoss # and StyleLoss we insert below. So we replace with out-of-place # ones here. (Not really understanding this...) layer = nn.ReLU(inplace=False) elif isinstance(layer, nn.MaxPool2d): name = 'pool_{}'.format(i) elif isinstance(layer, nn.BatchNorm2d): name = 'bn_{}'.format(i) else: raise RuntimeError('Unrecognized layer: {}'.format( layer.__class__.__name__)) model.add_module(name, layer) if name in content_layers: # add content loss target = model(content_img).detach() content_loss = ContentLoss(target) model.add_module('content_loss_{}'.format(i), content_loss) content_losses.append(content_loss) if name in style_layers: # add style loss target = model(style_img).detach() style_loss = StyleLoss(target) model.add_module('style_loss_{}'.format(i), style_loss) style_losses.append(style_loss) # trim off the layers after the last content and style losses for i in range(len(model) - 1, -1, -1): if isinstance(model[i], ContentLoss) or isinstance( model[i], StyleLoss): break model = model[:(i + 1)] return model, style_losses, content_losses
def __init__(self, args, dataloaders_xLabels_joint, dataloaders_single): super(train_inpainting_module_I, self).__init__(args) self._initialize_training() self.dataloaders_single = dataloaders_single self.dataloaders_xLabels_joint = dataloaders_xLabels_joint self.use_apex = False # use apex might cause style loss to be 0 self.mask_buffer = Mask_Buffer(500) self.attModule = _Attention_FullRes(input_nc=3, output_nc=1) # logits, no tanh() self.inpaintNet = _ResGenerator_Upsample(input_nc=3, output_nc=3) self.styleTranslator = _ResGenerator_Upsample(input_nc=3, output_nc=3) self.netD = DiscriminatorGlobalLocal(image_size=240) self.tau_min = 0.05 self.use_perceptual_loss = True self.p_vgg = VGG19() self.s_vgg = VGG19() self.perceptual_loss = PerceptualLoss(vgg19=self.p_vgg) self.style_loss = StyleLoss(vgg19=self.s_vgg) self.reconst_loss_weight = 1.0 self.perceptual_loss_weight = 1.0 self.style_loss_weight = 1.0 self.fake_loss_weight = 0.01 self.model_name = [ 'attModule', 'inpaintNet', 'styleTranslator', 'netD', 'p_vgg', 's_vgg' ] self.L1loss = nn.L1Loss() if self.isTrain: self.optim_inpaintNet = optim.Adam(self.inpaintNet.parameters(), lr=self.task_lr, betas=(0.5, 0.999)) self.optim_netD = optim.Adam(self.netD.parameters(), lr=self.task_lr, betas=(0.5, 0.999)) self.optim_name = ['optim_inpaintNet', 'optim_netD'] self._get_scheduler() self.loss_BCE = nn.BCEWithLogitsLoss() self._initialize_networks(['inpaintNet', 'netD']) # load the "best" style translator T (from step 2) preTrain_path = os.path.join(os.getcwd(), 'experiments', 'train_style_translator_T') self._load_models(model_list=['styleTranslator'], mode=480, isTrain=True, model_path=preTrain_path) print('Successfully loaded pre-trained {} model from {}'.format( 'styleTranslator', preTrain_path)) # load the "best" attention module A (from step 3) preTrain_path = os.path.join(os.getcwd(), 'experiments', 'train_initial_attention_module_A') self._load_models(model_list=['attModule'], mode=450, isTrain=True, model_path=preTrain_path) print('Successfully loaded pre-trained {} model from {}'.format( 'attModule', preTrain_path)) # apex can only be applied to CUDA models if self.use_apex: self._init_apex(Num_losses=2) self._check_parallel()
target = torch.autograd.Variable(torch.zeros_like(input_style), requires_grad=True) # perform forward pass of model to extract response for loss style_response = model.forward(input_style, extract_layers=arguments.style_layers) if output_content is not None: content_response = model.forward( output_content, extract_layers=arguments.content_layers) else: content_response = None # initialize loss style_loss = StyleLoss(style_response, input_map, output_map, arguments.style_layers, arguments.map_channel_weight, stride=1) if content_response is not None: content_loss = ContentLoss(content_response, arguments.content_layers) else: content_loss = None # setup optimizer optimizer = torch.optim.LBFGS([target], lr=1.0, history_size=100) # setup live plot if arguments.plot_interval is not None: _, _, height, width = input_style.size() live_plot = LivePlot(width, height)
def get_style_model_and_losses(cnn, target_texture_img, style_layers, structure_layers): histogram_layers = ['pool_4','pool_12'] input_channel = target_texture_img.size()[1] cnn = copy.deepcopy(cnn) output_channel, _, kernel_size, __ = cnn[0].weight.size() revised_cnn_weight = torch.zeros(output_channel, input_channel, cnn[0].kernel_size[0], cnn[0].kernel_size[1]).to(device) for ich in range(revised_cnn_weight.size()[1]): if ich < 3: revised_cnn_weight[:,ich,:,:] = cnn[0].weight.data[:,ich,:,:] else: t = torch.randint(3,(3,)).long() a = torch.rand(1).tolist()[0] b = (1-a)*torch.rand(1).tolist()[0] c = 1-a-b revised_cnn_weight[:,ich,:,:] = a*cnn[0].weight.data[:,t[0],:,:] +\ b*cnn[0].weight.data[:, t[1], :, :] +\ c*cnn[0].weight.data[:, t[2], :, :] new_first_layer = nn.Conv2d(input_channel,cnn[0].out_channels, kernel_size=cnn[0].kernel_size,padding=0, dilation=cnn[0].dilation,groups=cnn[0].groups) reflect_pad = nn.ReflectionPad2d((1,1,cnn[0].padding[0],cnn[0].padding[0])) new_first_layer.weight.data = revised_cnn_weight new_first_layer.bias.data = cnn[0].bias cnn[0] = new_first_layer # just in order to have an iterable access to or list of content/syle # losses structure_losses = [] histogram_losses = [] style_losses = [] # assuming that cnn is a nn.Sequential, so we make a new nn.Sequential # to put in modules that are supposed to be activated sequentially # model = nn.Sequential(normalization) model = nn.Sequential() model.add_module("reflect_pad",reflect_pad) tv_loss = TVloss() model.add_module("tv_loss", tv_loss) i = 0 # increment every time we see a conv for layer in cnn.children(): if isinstance(layer, nn.Conv2d): i += 1 name = 'conv_{}'.format(i) elif isinstance(layer, nn.ReLU): name = 'relu_{}'.format(i) # The in-place version doesn't play very nicely with the ContentLoss # and StyleLoss we insert below. So we replace with out-of-place # ones here. layer = nn.ReLU(inplace=False) elif isinstance(layer, nn.MaxPool2d): name = 'pool_{}'.format(i) layer = nn.AvgPool2d(kernel_size=2, stride=2) elif isinstance(layer, nn.BatchNorm2d): name = 'bn_{}'.format(i) else: raise RuntimeError('Unrecognized layer: {}'.format(layer.__class__.__name__)) model.add_module(name, layer) if name in style_layers: target_feature = model(target_texture_img).detach() style_loss = StyleLoss(target_feature) model.add_module("style_loss_{}".format(i), style_loss) style_losses.append(style_loss) if name in structure_layers: target_feature = model(target_texture_img).detach() structure_loss = StructureLoss(target_feature) model.add_module("structure_loss_{}".format(i), structure_loss) structure_losses.append(structure_loss) if name in histogram_layers: target_feature = model(target_texture_img).detach() histogram_loss = HistogramLoss(target_feature) model.add_module("histogram_loss_{}".format(i), histogram_loss) histogram_losses.append(histogram_loss) # now we trim off the layers after the last content and style losses for i in range(len(model) - 1, -1, -1): if isinstance(model[i], HistogramLoss): break model = model[:(i + 1)] return model, style_losses, structure_losses, tv_loss, histogram_losses
def sample_merging_optimization(init_guess, target_exemplar, upscaling_rate, img_res1,img_res2, kernel_sigma1,kernel_sigma2, num_optim_step, texture_weight, structure_weight, histogram_weight, texture_layers, structure_layers, optim_method, results_dir): global break_loop number_points_in_target_exemplar = target_exemplar.size()[0] neighborhood_size = 0.1 kernel_sigma_list = torch.linspace(kernel_sigma1, kernel_sigma2, 2) img_res_list = torch.linspace(img_res1, img_res2, 2) stopping_crit_list = torch.linspace(0.005, 0.005, 2) outerloop = 0 blue_noise = torch.from_numpy( generate_possion_dis(round(1 * (upscaling_rate ** 2) * number_points_in_target_exemplar), 0.0001, 0.9999)).to( device).float() init_guess_soft = torch.ones(init_guess.size(0),1).to(device) blue_noise_soft = torch.zeros(blue_noise.size(0),1).to(device) c = torch.cat((torch.ones(init_guess.size(0)),torch.zeros(blue_noise.size(0))),0) lr_list = [0.01, 0.002] lr_list_soft = [0.02, 0.02] lr_list_init_guess = [0.000, 5e-4] lr_list_init_guess_soft = [0.02, 0.02] run = [0] while outerloop < 2: stopping_crit = stopping_crit_list[outerloop].tolist() kernel_sigma = kernel_sigma_list[outerloop].tolist() img_res = img_res_list[outerloop].tolist() img_res_input = round(img_res * upscaling_rate) kernel_sigma_input = kernel_sigma/upscaling_rate from point2image import Point2Image as p2i target_p2i = p2i(2, 0, kernel_sigma=kernel_sigma, feature_sigma=0, res=img_res) input_p2i = SoftPoint2Image(2, 0, kernel_sigma=kernel_sigma_input, feature_sigma=0, res=img_res_input) if optim_method == 'LBFGS': optimizer = optim.LBFGS([blue_noise.requires_grad_()],lr=0.5) elif optim_method == 'Adam': optimizer = optim.Adam([ {'params': blue_noise.requires_grad_(), 'lr': lr_list[outerloop]}, {'params': blue_noise_soft.requires_grad_(), 'lr': lr_list_soft[outerloop]}, {'params': init_guess_soft.requires_grad_(), 'lr': lr_list_init_guess_soft[outerloop]}, {'params': init_guess.requires_grad_(), 'lr': lr_list_init_guess[outerloop]}]) print('step:', lr_list[outerloop]) target_texture_img = target_p2i(target_exemplar).repeat(1, 3, 1, 1).to(device) fstyle_loss = StyleLoss(target_texture_img) save_image(target_texture_img.squeeze()/ target_texture_img.max(), results_dir + '/target' + str(outerloop) + '.jpg') np.savetxt(results_dir + '/target' + str(outerloop) + '.txt', target_texture_img[0,0,:,:].cpu().data.numpy()) np.savetxt(results_dir + '/target_points' + str(outerloop) + '.txt', target_exemplar.cpu().data.numpy()) fig=plt.figure() plt.scatter(target_exemplar.data[:, 0], target_exemplar.data[:, 1]) plt.savefig(results_dir + '/scatter_target' + str(outerloop) + '.jpg') plt.close(fig) img_hist_loss = HistogramLoss(target_texture_img) print('Building the texture model..') model, texture_losses, structure_losses, _ , histogram_losses = get_style_model_and_losses(cnn, target_texture_img, texture_layers, structure_layers) log_losses = [] df_target = diff_func.apply(target_exemplar, neighborhood_size, 64, 2) df_target[0, 0, 28:36, 28:36] = 0 print('Optimizing..') break_loop = False inner_run = [0] while run[0] <= num_optim_step and not break_loop: def closure(): global break_loop blue_noise.data.clamp_(0, 1) init_guess.data.clamp_(0, 1) init_guess_soft.data.clamp_(0, 1) blue_noise_soft.data.clamp_(0, 1) input_pos = torch.cat((init_guess, blue_noise), 0) input_soft = torch.cat((init_guess_soft, blue_noise_soft),0) input = torch.cat((input_pos,input_soft),1) input_soft_points = input_p2i(input) input_soft_points.clamp_(min=target_texture_img.min()) optimizer.zero_grad() input_density_img = input_soft_points.repeat(1, 3, 1, 1) if run[0] == 0: save_image(input_density_img.squeeze() / input_density_img.max(), results_dir + '/init'+ str(outerloop) +'.jpg') np.savetxt(results_dir + '/init'+ str(outerloop) +'.txt', input_density_img[0,0,:,:].cpu().data.numpy()) np.savetxt(results_dir + '/init_points' + str(outerloop) + '.txt', input.cpu().data.numpy()) fig = plt.figure() plt.scatter(input.data[:, 0], input.data[:, 1],s=2,c=c) plt.savefig(results_dir + '/init_points' + str(outerloop) + '.jpg') plt.close(fig) model(input_density_img) img_hist_loss(input_density_img) fstyle_loss(input_density_img) img_hist_score = histogram_weight*img_hist_loss.loss texture_score = torch.zeros(1).to(device) structure_score = torch.zeros(1).to(device) histogram_score = torch.zeros(1).to(device) for tl in texture_losses: texture_score += texture_weight * tl.loss for sl in structure_losses: structure_score += structure_weight * sl.loss for hl in histogram_losses: histogram_score += histogram_weight*hl.loss.data.item() bound_score = torch.zeros(1).to(device) diff_score = torch.zeros(1).to(device) ftexture_score = fstyle_loss.loss loss = texture_score + structure_score + img_hist_score + histogram_score + bound_score + diff_score + ftexture_score #+ homo_score loss.backward() log_losses.append(loss) run[0] += 1 inner_run[0] += 1 if run[0] % 5 == 0: for param_group in optimizer.param_groups: print(param_group['lr']) print("run {}:".format(run)) print('Texture: {:4f}, FTexture: {:4f}, Structure: {:4f}, Image Hist : {:4f}, Histogram : {:4f}, Bound {:4f}, Diff {:4f}'.format( texture_score.item(), ftexture_score.item(), structure_score.item(),img_hist_score.item(),histogram_score.item(),bound_score.item(),diff_score.item())) save_density_img = input_density_img.clamp(min=target_texture_img.min(),max=target_texture_img.max()) save_image((save_density_img/save_density_img.max()), results_dir + '/out' + str(run[0]) + '_' + str(outerloop) + '.jpg') np.savetxt(results_dir + '/out' + str(run[0]) + '_' + str(outerloop) + '.txt', input_density_img[0,0,:,:].cpu().data.numpy()) np.savetxt(results_dir + '/out_points' + str(run[0]) + '_' + str(outerloop) + '.txt', input.cpu().data.numpy()) fig = plt.figure() plt.figure() plt.scatter(input.data[:, 0], input.data[:, 1],s=2,c=c) plt.savefig(results_dir + '/out_points' + str(run[0]) + '_' + str(outerloop) + '.jpg') plt.close('all') print('inner_run',inner_run) print('stopping_crit',stopping_crit) if inner_run[0] > 100: loss_init = log_losses[0] loss_pre = log_losses[inner_run[0]-100] loss_now = log_losses[inner_run[0]-1] decrease_perc = ((loss_pre-loss_now)/(loss_init-loss_now)).tolist()[0] print(decrease_perc) if decrease_perc < stopping_crit: print('converged') break_loop = True else: np.savetxt(results_dir + '/final_output' + '.txt', input.cpu().data.numpy()) np.savetxt(results_dir + '/final_output_density.txt', input_density_img[0, 0, :, :].cpu().data.numpy()) np.savetxt(results_dir + '/final_output' + '_' + str(outerloop) + '.txt', init_guess.cpu().data.numpy()) np.savetxt(results_dir + '/final_output_density' + '_' + str(outerloop) + '.txt', input_density_img[0, 0, :, :].cpu().data.numpy()) print() return loss optimizer.step(closure) outerloop +=1 blue_noise.data.clamp_(0, 1) init_guess.data.clamp_(0, 1) return blue_noise
select_tensor = CreateSpatialTensor().to(device) indices_generator = GetRandomIndices(device) try: content_image = load_image(args.content, args.max_resolution) style_image = load_image(args.style, args.max_resolution) except Exception as e: print(f'Error: {e}') sys.exit(1) # Calculating depth of laplacian pyramid pyr_depth = int(math.log2(content_image.shape[-1])) - 6 + 1 image_pyramid = LaplacianPyramid(pyr_depth).to(device) origin_image_pyramid = image_pyramid(content_image.to(device)) criteria = StyleLoss() alpha = args.alpha content_pyramid = None for j, con_image in enumerate(origin_image_pyramid[1:]): batch, channel, height, width = con_image.shape print(f'Image size: ({height}, {width})') org_con_image = F.interpolate(content_image, (height, width), mode='bilinear', align_corners=False).to(device) st_image = F.interpolate(style_image, (height, width), mode='bilinear', align_corners=False).to(device) with torch.no_grad():
def get_style_model_and_losses(cnn, style_img, content_img, device=torch.device('cpu')): """ We need to add our content loss and style loss layers immediately after the convolution layer they are detecting. To do this we must create a new Sequential module that has content loss and style loss modules correctly inserted. """ # desired depth layers to compute style/content losses: content_layers = ['conv_4'] style_layers = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5'] cnn = copy.deepcopy(cnn) normalization = Normalization().to(device) content_losses, style_losses = [], [] model = nn.Sequential(normalization) n_conv = 0 for layer in cnn.children(): if isinstance(layer, nn.Conv2d): n_conv += 1 name = 'conv_{}'.format(n_conv) elif isinstance(layer, nn.ReLU): name = 'relu_{}'.format(n_conv) layer = nn.ReLU(inplace=False) elif isinstance(layer, nn.MaxPool2d): name = 'pool_{}'.format(n_conv) elif isinstance(layer, nn.BatchNorm2d): name = 'bn_{}'.format(n_conv) else: raise RuntimeError('Unrecognized layer:{}'.format( layer.__class__.__name__)) model.add_module(name, layer) if name in content_layers: # add content loss target_feature = model(content_img).detach() content_loss = ContentLoss(target_feature) model.add_module('content_loss_{}'.format(n_conv), content_loss) content_losses.append(content_loss) if name in style_layers: # add style loss target_feature = model(style_img).detach() style_loss = StyleLoss(target_feature) model.add_module('style_loss_{}'.format(n_conv), style_loss) style_losses.append(style_loss) for i in range(len(model) - 1, -1, -1): if isinstance(model[i], ContentLoss) or isinstance( model[i], StyleLoss): break model = model[:(i + 1)] return model, style_losses, content_losses
def hierarchical_end_to_end_optimization_sample_position( init_guess, target_exemplar, upscaling_rate, img_res1, img_res2, kernel_sigma1, kernel_sigma2, num_optim_step, texture_weight, structure_weight, histogram_weight, image_histogram_weight, texture_layers, structure_layers, optim_method, results_dir): global break_loop kernel_sigma_list = torch.linspace(kernel_sigma1, kernel_sigma2, 3) img_res_list = torch.linspace(img_res1, img_res2, 3) stopping_crit_list = torch.linspace(0.01, 0.01, 3) outerloop = 0 lr_list = [0.02, 0.01, 0.01, 0.002, 0.002, 0.002, 0.002, 0.002] run = [0] while outerloop < 3: stopping_crit = stopping_crit_list[outerloop].tolist() kernel_sigma = kernel_sigma_list[outerloop].tolist() img_res = img_res_list[outerloop].tolist() img_res_input = round(img_res * upscaling_rate) kernel_sigma_input = kernel_sigma / upscaling_rate from point2image import Point2Image target_p2i = Point2Image(2, 0, kernel_sigma=kernel_sigma, feature_sigma=0, res=img_res) input_p2i = Point2Image(2, 0, kernel_sigma=kernel_sigma_input, feature_sigma=0, res=img_res_input) if optim_method == 'LBFGS': optimizer = optim.LBFGS([init_guess.requires_grad_()], lr=0.5) elif optim_method == 'Adam': optimizer = optim.Adam([init_guess.requires_grad_()], lr=lr_list[outerloop]) print('step:', lr_list[outerloop]) target_texture_img = target_p2i(target_exemplar).repeat(1, 3, 1, 1).to(device) fstyle_loss = StyleLoss(target_texture_img) save_image(target_texture_img.squeeze() / target_texture_img.max(), results_dir + '/target' + str(outerloop) + '.jpg') np.savetxt(results_dir + '/target' + str(outerloop) + '.txt', target_texture_img[0, 0, :, :].cpu().data.numpy()) np.savetxt(results_dir + '/target_points' + str(outerloop) + '.txt', target_exemplar.cpu().data.numpy()) fig = plt.figure() plt.scatter(target_exemplar.data[:, 0], target_exemplar.data[:, 1]) plt.savefig(results_dir + '/scatter_target' + str(outerloop) + '.jpg') plt.close(fig) img_hist_loss = HistogramLoss(target_texture_img) print('Building the texture model..') model, texture_losses, structure_losses, _, histogram_losses = get_style_model_and_losses( cnn, target_texture_img, texture_layers, structure_layers) log_losses = [] print('Optimizing..') break_loop = False inner_run = [0] while run[0] <= num_optim_step and not break_loop: def closure(): global break_loop init_guess.data.clamp_(0, 1) input_soft_points = input_p2i(init_guess) input_soft_points.clamp(min=target_texture_img.min(), max=target_texture_img.max()) optimizer.zero_grad() input_density_img = input_soft_points.repeat(1, 3, 1, 1) img_hist_loss(input_density_img) if run[0] == 0: save_image( input_density_img.squeeze() / input_density_img.max(), results_dir + '/init' + str(outerloop) + '.jpg') np.savetxt( results_dir + '/init' + str(outerloop) + '.txt', input_density_img[0, 0, :, :].cpu().data.numpy()) np.savetxt( results_dir + '/init_points' + str(outerloop) + '.txt', init_guess.cpu().data.numpy()) fig = plt.figure() plt.scatter(init_guess.data[:, 0], init_guess.data[:, 1]) plt.savefig(results_dir + '/init_points' + str(outerloop) + '.jpg') plt.close(fig) model(input_density_img) texture_score = torch.zeros(1).to(device) structure_score = torch.zeros(1).to(device) histogram_score = torch.zeros(1).to(device) img_hist_score = torch.zeros(1).to(device) for tl in texture_losses: texture_score += texture_weight * tl.loss for sl in structure_losses: structure_score += structure_weight * sl.loss fstyle_loss(input_density_img) ftexture_score = fstyle_loss.loss loss = texture_score + structure_score + histogram_score + img_hist_score + ftexture_score #+ homo_score loss.backward() log_losses.append(loss) run[0] += 1 inner_run[0] += 1 if run[0] % 5 == 0: for param_group in optimizer.param_groups: print(param_group['lr']) print("run {}:".format(run)) print( 'Texture Loss : {:4f}, FTexture_Loss: {:4f}, Structure Loss : {:4f}, Histogram Loss : {:4f}, Image Histogram Loss : {:4f}' .format(texture_score.item(), ftexture_score.item(), structure_score.item(), histogram_score.item(), img_hist_score.item())) save_image((input_density_img / input_density_img.max()), results_dir + '/out' + str(run[0]) + '_' + str(outerloop) + '.jpg') np.savetxt( results_dir + '/out' + str(run[0]) + '_' + str(outerloop) + '.txt', input_density_img[0, 0, :, :].cpu().data.numpy()) np.savetxt( results_dir + '/out_points' + str(run[0]) + '_' + str(outerloop) + '.txt', init_guess.cpu().data.numpy()) fig = plt.figure() plt.figure() plt.scatter(init_guess.data[:, 0], init_guess.data[:, 1], s=2) plt.savefig(results_dir + '/out_points' + str(run[0]) + '_' + str(outerloop) + '.jpg') plt.close('all') print('inner_run', inner_run) print('stopping_crit', stopping_crit) if inner_run[0] > 100: loss_init = log_losses[0] loss_pre = log_losses[inner_run[0] - 100] loss_now = log_losses[inner_run[0] - 1] decrease_perc = ((loss_pre - loss_now) / (loss_init - loss_now)).tolist()[0] print(decrease_perc) if decrease_perc < stopping_crit: print('converged') break_loop = True else: np.savetxt( results_dir + '/final_output' + '_' + str(outerloop) + '.txt', init_guess.cpu().data.numpy()) np.savetxt( results_dir + '/final_output_density' + '_' + str(outerloop) + '.txt', input_density_img[0, 0, :, :].cpu().data.numpy()) np.savetxt(results_dir + '/final_output' + '.txt', init_guess.cpu().data.numpy()) np.savetxt( results_dir + '/final_output_density.txt', input_density_img[0, 0, :, :].cpu().data.numpy()) print() with open( results_dir + '/log_losses_' + str(outerloop) + '.txt', "w") as f: for s in log_losses: f.write(str(s.tolist()[0]) + "\n") return loss optimizer.step(closure) outerloop += 1 init_guess.data.clamp_(0, 1) return init_guess
# build the neural network to optimize for layer in list(vgg19): if isinstance(layer, nn.Conv2d): conv = 'conv_' + str(i) model.add_module(conv, layer) if conv in content_layers: target = model.forward(photo).clone() cl = ContentLoss(args.w_content, target) model.add_module('content_loss_' + str(i), cl) cl_list.append(cl) if conv in style_layers: fm = model.forward(art).clone() target_fm = gm.forward(fm) sl = StyleLoss(args.w_style, target_fm) model.add_module('style_loss_' + str(i), sl) sl_list.append(sl) i += 1 if isinstance(layer, nn.ReLU): relu = 'relu_' + str(i) model.add_module(relu, layer) if isinstance(layer, nn.MaxPool2d): mpool = 'maxpool_' + str(i) model.add_module(mpool, layer) print(model) # setup inputs and optimizers
def get_style_model_and_losses(cnn, normalization_mean, normalization_std, style_img, content_img, content_layers, style_layers, device): cnn = copy.deepcopy(cnn) # normalization module normalization = Normalization(normalization_mean, normalization_std).to(device) # just in order to have an iterable access to or list of content/syle # losses content_losses = [] style_losses = [] # assuming that cnn is a nn.Sequential, so we make a new nn.Sequential # to put in modules that are supposed to be activated sequentially ''' (1)给自定义的model建立了一个Sequential,万一以后有其他用处就方便点,如果没有这个nn.Sequential(), 都不知道下面能不能用那个model.add_module(name, layer) ,为了方便以后自己也得有nn.Sequential();有个问题就是我怎么自定义呢? 像该文件代码那样不建立类情况下,自定义出来一个像nn.features()和nn.classifier()一样的东西。下面代码亲自试验有效,可以任意订制模型: model = nn.Sequential() y = nn.Sequential( nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)), nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))) y.add_module('haha', nn.ReLU()) model.add_module('features', y) z = nn.Sequential(nn.Dropout(p=0.5), nn.Dropout(p=0.5)) model.add_module('classifier', z) print(model) print(model.features) 至于怎么运行,就像下面代码一样,定义好结构,然后不用定义forward,直接model(input),就执行了,记得把model各个块(各层) 的参数放进优化器即可,其他什么反向传播等等都和调用模型时语法流程一样的。 (2)建立类时: class LeNet(nn.Module): def __init__(self): super(LeNet, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 6, 5), nn.ReLU(), nn.MaxPool2d(2, 2), nn.Conv2d(6, 16, 5), nn.ReLU(), nn.MaxPool2d(2, 2) ) # 由于调整shape并不是一个class层, # 所以在涉及这种操作(非nn.Module操作)需要拆分为多个模型 self.classifiter = nn.Sequential( nn.Linear(16*5*5, 120), nn.ReLU(), nn.Linear(120, 84), nn.ReLU(), nn.Linear(84, 10) ) def forward(self, x): x = self.features(x) x = x.view(-1, 16*5*5) x = self.classifiter(x) return x ''' model = nn.Sequential(normalization) #normalization才是model[0] '''相当于我自己要从头开始一块块(一个积木,一层,一个零件)建立一个模型了''' ''' 思想是在已有模型基础上进行分拆,添加,组合;而到底怎么做,要根据你利用的已有模型的结构来分析; 比如VGG模型就是Conv,Relu,BN,Pool交替循环,所以才有了下面为什么这么写。 ''' i = 0 # increment every time we see a conv for layer in cnn.children(): if isinstance(layer, nn.Conv2d): i += 1 name = 'conv_%d' % i elif isinstance(layer, nn.ReLU): name = 'relu_%d' % i # The in-place version doesn't play very nicely with the ContentLoss # and StyleLoss we insert below. So we replace with out-of-place # ones here. layer = nn.ReLU(inplace=False) elif isinstance(layer, nn.BatchNorm2d): name = 'bn_%d' % i elif isinstance(layer, nn.MaxPool2d): name = 'pool_%d' % i else: raise RuntimeError('Unrecognized layer: %s' % (layer.__class__.__name__)) model.add_module(name, layer) # 把layer添加到model并取名为name if name in content_layers: # add content loss: target = model(content_img).detach( ) # 跑一下模型,现在定义了多少多长跑多长,就可以得到中间的C x features map content_loss = ContentLoss( target) # 对ContentLoss这个模型类进行实例化,但还没有赋值进行forward model.add_module("content_loss_%d" % i, content_loss) # 将这个实例添加进模型 content_losses.append(content_loss) # 实例添加进列表 if name in style_layers: # add style loss: target_feature = model(style_img).detach() style_loss = StyleLoss(target_feature) model.add_module("style_loss_%d" % i, style_loss) style_losses.append(style_loss) # now we trim off the layers after the last content and style losses for i in range(len(model) - 1, -1, -1): # 从最后一层遍历到第一层 if isinstance(model[i], ContentLoss) or isinstance( model[i], StyleLoss): break model = model[:(i + 1)] # 这个模型后面用不到的几层既可以截掉啦 return model, style_losses, content_losses