168 lines
7.5 KiB
Python
168 lines
7.5 KiB
Python
import torch
|
|
from torch.optim import AdamW, LBFGS
|
|
from torchvision import models
|
|
import torch.nn as nn
|
|
import argparse
|
|
from collections import namedtuple
|
|
|
|
class Vgg16(torch.nn.Module):
|
|
"""Only those layers are exposed which have already proven to work nicely."""
|
|
def __init__(self, requires_grad=False, show_progress=False):
|
|
super().__init__()
|
|
vgg_pretrained_features = models.vgg16(pretrained=True, progress=show_progress).features
|
|
self.layer_names = ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3']
|
|
self.content_feature_maps_index = 1 # relu2_2
|
|
self.style_feature_maps_indices = list(range(len(self.layer_names))) # all layers used for style representation
|
|
|
|
self.slice1 = torch.nn.Sequential()
|
|
self.slice2 = torch.nn.Sequential()
|
|
self.slice3 = torch.nn.Sequential()
|
|
self.slice4 = torch.nn.Sequential()
|
|
for x in range(4):
|
|
self.slice1.add_module(str(x), vgg_pretrained_features[x])
|
|
for x in range(4, 9):
|
|
self.slice2.add_module(str(x), vgg_pretrained_features[x])
|
|
for x in range(9, 16):
|
|
self.slice3.add_module(str(x), vgg_pretrained_features[x])
|
|
for x in range(16, 23):
|
|
self.slice4.add_module(str(x), vgg_pretrained_features[x])
|
|
if not requires_grad:
|
|
for param in self.parameters():
|
|
param.requires_grad = False
|
|
|
|
def forward(self, x):
|
|
x = self.slice1(x)
|
|
relu1_2 = x
|
|
x = self.slice2(x)
|
|
relu2_2 = x
|
|
x = self.slice3(x)
|
|
relu3_3 = x
|
|
x = self.slice4(x)
|
|
relu4_3 = x
|
|
vgg_outputs = namedtuple("VggOutputs", self.layer_names)
|
|
out = vgg_outputs(relu1_2, relu2_2, relu3_3, relu4_3)
|
|
return out
|
|
|
|
|
|
def prepare_model(device):
|
|
model = Vgg16()
|
|
|
|
content_feature_maps_index = model.content_feature_maps_index
|
|
style_feature_maps_indices = model.style_feature_maps_indices
|
|
layer_names = model.layer_names
|
|
|
|
content_fms_index_name = (content_feature_maps_index, layer_names[content_feature_maps_index])
|
|
style_fms_indices_names = (style_feature_maps_indices, layer_names)
|
|
return model.to(device).eval(), content_fms_index_name, style_fms_indices_names
|
|
|
|
|
|
def gram_matrix(x, should_normalize=True):
|
|
(b, ch, h, w) = x.size()
|
|
features = x.view(b, ch, w * h)
|
|
features_t = features.transpose(1, 2)
|
|
gram = features.bmm(features_t)
|
|
if should_normalize:
|
|
gram /= ch * h * w
|
|
return gram
|
|
|
|
|
|
def total_variation(y):
|
|
return torch.sum(torch.abs(y[:, :, :, :-1] - y[:, :, :, 1:])) + \
|
|
torch.sum(torch.abs(y[:, :, :-1, :] - y[:, :, 1:, :]))
|
|
|
|
|
|
def build_loss(neural_net, optimizing_img, target_representations, content_feature_maps_index, style_feature_maps_indices, config):
|
|
target_content_representation = target_representations[0]
|
|
target_style_representation = target_representations[1]
|
|
|
|
current_set_of_feature_maps = neural_net(optimizing_img)
|
|
|
|
current_content_representation = current_set_of_feature_maps[content_feature_maps_index].squeeze(axis=0)
|
|
content_loss = torch.nn.MSELoss(reduction='mean')(target_content_representation, current_content_representation)
|
|
|
|
style_loss = 0.0
|
|
current_style_representation = [gram_matrix(x) for cnt, x in enumerate(current_set_of_feature_maps) if cnt in style_feature_maps_indices]
|
|
for gram_gt, gram_hat in zip(target_style_representation, current_style_representation):
|
|
style_loss += torch.nn.MSELoss(reduction='sum')(gram_gt[0], gram_hat[0])
|
|
style_loss /= len(target_style_representation)
|
|
|
|
tv_loss = total_variation(optimizing_img)
|
|
|
|
total_loss = config['content_weight'] * content_loss + config['style_weight'] * style_loss + config['tv_weight'] * tv_loss
|
|
|
|
return total_loss, content_loss, style_loss, tv_loss
|
|
|
|
def make_tuning_step(neural_net, optimizer, target_representations, content_feature_maps_index, style_feature_maps_indices, config):
|
|
def tuning_step(optimizing_img):
|
|
optimizer.zero_grad()
|
|
total_loss, content_loss, style_loss, tv_loss = build_loss(neural_net, optimizing_img, target_representations, content_feature_maps_index, style_feature_maps_indices, config)
|
|
total_loss.backward()
|
|
optimizer.step()
|
|
return total_loss, content_loss, style_loss, tv_loss
|
|
return tuning_step
|
|
|
|
def neural_style_transfer_from_tensors(content_img_tensor, style_img_tensor, config, num_steps, learning_rate):
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
content_img = content_img_tensor.to(device)
|
|
style_img = style_img_tensor.to(device)
|
|
|
|
optimizing_img = content_img.clone().requires_grad_(True)
|
|
|
|
neural_net, content_feature_maps_index_name, style_feature_maps_indices_names = prepare_model(device)
|
|
|
|
content_img_set_of_feature_maps = neural_net(content_img)
|
|
style_img_set_of_feature_maps = neural_net(style_img)
|
|
|
|
target_content_representation = content_img_set_of_feature_maps[content_feature_maps_index_name[0]].squeeze(axis=0)
|
|
target_style_representation = [gram_matrix(x) for cnt, x in enumerate(style_img_set_of_feature_maps) if cnt in style_feature_maps_indices_names[0]]
|
|
target_representations = [target_content_representation, target_style_representation]
|
|
|
|
optimizer = LBFGS((optimizing_img,), max_iter=num_steps, lr=learning_rate, line_search_fn='strong_wolfe')
|
|
cnt = 0
|
|
|
|
def closure():
|
|
nonlocal cnt
|
|
if torch.is_grad_enabled():
|
|
optimizer.zero_grad()
|
|
total_loss, content_loss, style_loss, tv_loss = build_loss(neural_net, optimizing_img, target_representations, content_feature_maps_index_name[0], style_feature_maps_indices_names[0], config)
|
|
if total_loss.requires_grad:
|
|
total_loss.backward()
|
|
if cnt%100==0:
|
|
with torch.no_grad():
|
|
print(f'L-BFGS | iteration: {cnt:03}, total loss={total_loss.item():12.4f}, content_loss={config["content_weight"] * content_loss.item():12.4f}, style loss={config["style_weight"] * style_loss.item():12.4f}, tv loss={config["tv_weight"] * tv_loss.item():12.4f}')
|
|
|
|
cnt += 1
|
|
return total_loss
|
|
|
|
optimizer.step(closure)
|
|
|
|
return optimizing_img.detach()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--temp_save_style_img", type=str)
|
|
parser.add_argument("--temp_save_content_img", type=str)
|
|
parser.add_argument("--temp_load_final_img", type=str)
|
|
|
|
parser.add_argument("--content_weight", type=float, help="weight factor for content loss", default=1e5)
|
|
parser.add_argument("--style_weight", type=float, help="weight factor for style loss", default=3e4)
|
|
parser.add_argument("--tv_weight", type=float, help="weight factor for total variation loss", default=1e0)
|
|
parser.add_argument("--learning_rate", type=float, help="learning_rate", default=1e1)
|
|
parser.add_argument("--num_steps", type=int, help="number of training steps", default=100)
|
|
|
|
args = parser.parse_args()
|
|
config = {
|
|
"content_weight": args.content_weight,
|
|
"style_weight": args.style_weight,
|
|
"tv_weight": args.tv_weight,
|
|
}
|
|
style_img = torch.load(args.temp_save_style_img)
|
|
content_img = torch.load(args.temp_save_content_img)
|
|
final_img = neural_style_transfer_from_tensors(content_img, style_img, config, args.num_steps, args.learning_rate)
|
|
print("nst done")
|
|
|
|
torch.save(final_img, args.temp_load_final_img)
|
|
|