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