Style Transfer 记录

经典文章xxx，理论就不介绍了，根据一个content图像，和一个style图像，可以把style图像的style迁移到content图像上。
在代码上有一个跟之前不同的地方，就是这里需要不断优化的变量是这张图像，vgg只是用来提取特征，不需要反传。具体做的时候，把三张图(content 图， style图，和我们希望生成的target图)都通过vgg，提取中间某些层计算出feature，这些feature 之间会计算一个content loss，使得target图和content图内容接近，同时计算一个style loss，使得target图和style图的style接近。

一般为了加速迭代，target会先用content图初始化。

不废话，上代码

import os
os.chdir(os.path.dirname(__file__))
from torchvision import models
from torchvision import transforms
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import argparse
from PIL import Image
from torch.utils.tensorboard import SummaryWritersample_dir = 'samples_style_transfer'
if not os.path.exists(sample_dir):os.makedirs(sample_dir, exist_ok=True)writer = SummaryWriter(sample_dir)device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')def load_image(image_path, transform=None, max_size=None, shape=None):image = Image.open(image_path)if max_size:scale = max_size / max(image.size)size = np.array(image.size) * scaleimage = image.resize(size.astype(int), Image.ANTIALIAS)if shape:image = image.resize(shape, Image.LANCZOS)if transform:image = transform(image).unsqueeze(0)return image.to(device)class VGGNet(nn.Module):def __init__(self):super(VGGNet, self).__init__()self.select = ['0', '5', '10', '19', '28'] self.vgg = models.vgg19(pretrained=True).featuresdef forward(self, x):features = []for name, layer in self.vgg._modules.items():x = layer(x)if name in self.select:features.append(x)return featuresdef main(config):T = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))])content = load_image(config.content, T, max_size=config.max_size)style = load_image(config.style, T, shape=[content.size(2), content.size(3)])target = content.clone().requires_grad_(True)optimizer = torch.optim.Adam([target], lr=config.lr, betas=[0.5, 0.999])vgg = VGGNet().to(device).eval()for epoch in range(config.total_step):target_feature = vgg(target)content_feature = vgg(content)style_feature = vgg(style)style_loss = 0content_loss = 0for f1, f2, f3 in zip(target_feature, content_feature, style_feature):content_loss += torch.mean((f1-f2)**2)_,c,h,w = f1.size()f1 = f1.view(c, h*w)f3 = f3.view(c, h*w)# gram matrixf1 = torch.mm(f1, f1.t())f3 = torch.mm(f3, f3.t())# style lossstyle_loss += torch.mean((f1-f3)**2) / (c*h*w)loss = content_loss + config.style_weight * style_lossoptimizer.zero_grad()loss.backward()optimizer.step()writer.add_scalar('loss', loss.item(), global_step=epoch)writer.add_scalar('content_loss', content_loss.item(), global_step=epoch)writer.add_scalar('style_loss', style_loss.item(), global_step=epoch)if (epoch+1) % config.log_step == 0:print('Epoch [{}/{}], Loss: {:.4f}, Content loss: {:.4f}, Style loss: {:.4f}'.\format(epoch, config.total_step, loss.item(), content_loss.item(), style_loss.item()))if (epoch+1) % config.sample_step == 0:denorm = transforms.Normalize(mean=(-2.12, -2.04, -1.80), std=(4.37, 4.46, 4.44))img = target.clone().squeeze()img = denorm(img).clamp_(0, 1)writer.add_image('img', img, global_step=epoch)if __name__ == "__main__":parser = argparse.ArgumentParser()parser.add_argument('--content', type=str, default='data/content.png')parser.add_argument('--style', type=str, default='data/style.png')parser.add_argument('--max_size', type=int, default=400)parser.add_argument('--total_step', type=int, default=2000)parser.add_argument('--log_step', type=int, default=10)parser.add_argument('--sample_step', type=int, default=500)parser.add_argument('--style_weight', type=float, default=100)parser.add_argument('--lr', type=float, default=0.003)config = parser.parse_args()config.total_step = 20000config.sample_step = 100print(config)main(config)

content图
style图
接近4000次迭代
经过 20000次迭代，能够看到风格越来越接近了。

但是这种style transfer有一个缺点，就是得针对一张图像进行不断迭代，能不能来了一张新图像，送进去后很快就能得到新的图像呢？

当然有，那就是fast neural style transfer

Fast Style Transfer 记录

相关原理可以参考 https://blog.csdn.net/qq_33590958/article/details/96122789
相当于设计了一个transform 的网络，专门做风格转换，然后继续保留预训练好的vgg网络用来提取特征做loss，整体loss还是参考了原来的文章，有了这个transform网络，只需要训练好这个网络，新的图像输进去，一个前向过程就输出了转换后的图像，速度很快。不过也有个缺陷，那就是每个新的style图，都需要重新训练网络(我的理解是这样，有错的话欢迎指出)。
在这里插入图片描述
具体transform的网络结构如下，是全连接网络，大致分三个阶段，下采样-> 残差模块->上采样模块

代码是参考了pytorch的example代码，如下

import os
# os.chdir(os.path.dirname(__file__))
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torchvision
from torchvision import transforms
from torchvision import datasets
from torchvision import models
from torch.utils.tensorboard import SummaryWriter
import numpy as np
from PIL import Image
import argparsesample_dir = 'samples_fast_style_transfer'
if not os.path.exists(sample_dir):os.makedirs(sample_dir, exist_ok=True)writer = SummaryWriter(sample_dir)device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
np.random.seed(0)
torch.manual_seed(0)def load_image(filename, size=None, scale=None):img = Image.open(filename).convert('RGB')if size is not None:img = img.resize((size, size), Image.ANTIALIAS)elif scale is not None:size = (int(img.size[0] / scale), int(img.size[1] / scale))img = img.resize(size, Image.ANTIALIAS)return imgdef save_image(filename, data):img = data.clone.clamp(0, 255).numpy()img = img.transpose(1,2,0).astype('uint8')img = Image.fromarray(img)img.save(filename)def gram_matrix(y):b,ch,h,w = y.size()features = y.view(b,ch,h*w)features_t = features.transpose(1,2)gram = features.bmm(features_t)/(ch*h*w)return gramdef normalize_batch(batch):# normalize using imagenet mean and stdmean = batch.new_tensor([0.485, 0.456, 0.406]).view(-1, 1, 1)std = batch.new_tensor([0.229, 0.224, 0.225]).view(-1, 1, 1)batch = batch.div_(255.0)return (batch - mean) / stddef save_checkpoint(model, epochs):model.eval().cpu()save_model_filename = "epoch_" + str(epochs) + ".model"save_model_path = os.path.join(sample_dir, save_model_filename)torch.save(model.state_dict(), save_model_path)print("\nDone, trained model saved at", save_model_path)class ConvLayer(torch.nn.Module):def __init__(self, in_channels, out_channels, kernel_size, stride):super(ConvLayer, self).__init__()reflection_padding = kernel_size // 2self.reflection_pad = torch.nn.ReflectionPad2d(reflection_padding)self.conv2d = torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride)def forward(self, x):out = self.reflection_pad(x)out = self.conv2d(out)return outclass ResidualBlock(torch.nn.Module):"""ResidualBlockintroduced in: https://arxiv.org/abs/1512.03385recommended architecture: http://torch.ch/blog/2016/02/04/resnets.html"""def __init__(self, channels):super(ResidualBlock, self).__init__()self.conv1 = ConvLayer(channels, channels, kernel_size=3, stride=1)self.in1 = torch.nn.InstanceNorm2d(channels, affine=True)self.conv2 = ConvLayer(channels, channels, kernel_size=3, stride=1)self.in2 = torch.nn.InstanceNorm2d(channels, affine=True)self.relu = torch.nn.ReLU()def forward(self, x):residual = xout = self.relu(self.in1(self.conv1(x)))out = self.in2(self.conv2(out))out = out + residualreturn outclass UpsampleConvLayer(torch.nn.Module):"""UpsampleConvLayerUpsamples the input and then does a convolution. This method gives better resultscompared to ConvTranspose2d.ref: http://distill.pub/2016/deconv-checkerboard/"""def __init__(self, in_channels, out_channels, kernel_size, stride, upsample=None):super(UpsampleConvLayer, self).__init__()self.upsample = upsamplereflection_padding = kernel_size // 2self.reflection_pad = torch.nn.ReflectionPad2d(reflection_padding) # 可以重点关注这个padding方式,参考https://blog.csdn.net/LionZYT/article/details/120181586self.conv2d = torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride)def forward(self, x):x_in = xif self.upsample:x_in = torch.nn.functional.interpolate(x_in, mode='nearest', scale_factor=self.upsample)out = self.reflection_pad(x_in)out = self.conv2d(out)return outclass TransformerNet(torch.nn.Module):def __init__(self):super(TransformerNet, self).__init__()# Initial convolution layersself.conv1 = ConvLayer(3, 32, kernel_size=9, stride=1)self.in1 = torch.nn.InstanceNorm2d(32, affine=True)self.conv2 = ConvLayer(32, 64, kernel_size=3, stride=2)self.in2 = torch.nn.InstanceNorm2d(64, affine=True)self.conv3 = ConvLayer(64, 128, kernel_size=3, stride=2)self.in3 = torch.nn.InstanceNorm2d(128, affine=True)# Residual layersself.res1 = ResidualBlock(128)self.res2 = ResidualBlock(128)self.res3 = ResidualBlock(128)self.res4 = ResidualBlock(128)self.res5 = ResidualBlock(128)# Upsampling Layersself.deconv1 = UpsampleConvLayer(128, 64, kernel_size=3, stride=1, upsample=2)self.in4 = torch.nn.InstanceNorm2d(64, affine=True)self.deconv2 = UpsampleConvLayer(64, 32, kernel_size=3, stride=1, upsample=2)self.in5 = torch.nn.InstanceNorm2d(32, affine=True)self.deconv3 = ConvLayer(32, 3, kernel_size=9, stride=1)# Non-linearitiesself.relu = torch.nn.ReLU()def forward(self, X):y = self.relu(self.in1(self.conv1(X)))y = self.relu(self.in2(self.conv2(y)))y = self.relu(self.in3(self.conv3(y)))y = self.res1(y)y = self.res2(y)y = self.res3(y)y = self.res4(y)y = self.res5(y)y = self.relu(self.in4(self.deconv1(y)))y = self.relu(self.in5(self.deconv2(y)))y = self.deconv3(y)return yfrom collections import namedtuple
class Vgg16(nn.Module):def __init__(self, required_grad=False):super(Vgg16, self).__init__()vgg_pretrained_features = models.vgg16(pretrained=True).featuresself.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 required_grad:for param in self.parameters():param.requires_grad_(False)def forward(self, X):h = self.slice1(X)h_relu1_2 = hh = self.slice2(h)h_relu2_2 = hh = self.slice3(h)h_relu3_3 = hh = self.slice4(h)h_relu4_3 = hvgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3'])out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3)return outdef train():epochs = 10batch_size = 4image_size = 256learning_rate = 1e-3style_weight = 1e6content_weight = 1e1log_step = 10dataset_path = 'data/fast_neural_style/dataset'style_image = 'data/fast_neural_style/style-images/mosaic.jpg'content_image = 'data/fast_neural_style/content-images/amber.jpg'T = transforms.Compose([transforms.Resize(image_size),transforms.CenterCrop(image_size),transforms.ToTensor(),transforms.Lambda(lambda x : x.mul(255))])train_dataset = datasets.ImageFolder(dataset_path, T)train_loader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True, shuffle=True)transformer = TransformerNet().to(device)optimizer = torch.optim.Adam(transformer.parameters(), lr=learning_rate)mse_loss = torch.nn.MSELoss()vgg = Vgg16(required_grad=False).to(device)# load style imagestyle_T = transforms.Compose([transforms.ToTensor(),transforms.Lambda(lambda x :x.mul(255))])style = load_image(style_image, size = image_size)style = style_T(style)style = style.repeat(batch_size, 1,1,1).to(device)features_style = vgg(normalize_batch(style))gram_style = [gram_matrix(y) for y in features_style]cnt = 0# load content imagecontent_image = load_image(content_image)content_transform = transforms.Compose([transforms.ToTensor(),transforms.Lambda(lambda x: x.mul(255))])content_image = content_transform(content_image)content_image = content_image.unsqueeze(0).to(device)for epoch in range(epochs):transformer.train()for batchid, (x, _) in enumerate(train_loader):x = x.to(device)y = transformer(x)# 这里是为了让数据的均值和方差符合预训练模型的分布x = normalize_batch(x)y = normalize_batch(y.mul(255))features_x = vgg(x)features_y = vgg(y)content_loss = mse_loss(features_y.relu2_2, features_x.relu2_2) * content_weightstyle_loss = 0for ft_y, gm_s in zip(features_y, gram_style):gm_y = gram_matrix(ft_y)style_loss += mse_loss(gm_y, gm_s)style_loss = style_loss * style_weighttotal_loss = content_loss  + style_loss optimizer.zero_grad()total_loss.backward()optimizer.step()cnt += 1if cnt % log_step == 0:print('Epoch [{}/{}], Step [{}], Loss: {:.4f}, Content loss: {:.4f}, Style loss: {:.4f}'.\format(epoch, epochs, cnt, total_loss.item(), content_loss.item(), style_loss.item()))writer.add_scalar('loss', total_loss.item(), global_step=cnt)writer.add_scalar('content_loss', content_loss.item(), global_step=cnt)writer.add_scalar('style_loss', style_loss.item(), global_step=cnt)if cnt % 100 == 0:img = eval(content_image, transformer)writer.add_image('target_images', img, global_step=cnt, dataformats='CHW')save_checkpoint(transformer, epoch)def eval(content_image, transformer):transformer.eval()output_image = transformer(content_image).cpu()transformer.train()mean = output_image.new_tensor([0.485, 0.456, 0.406]).view(-1, 1, 1)std = output_image.new_tensor([0.229, 0.224, 0.225]).view(-1, 1, 1)output_image = output_image * std + meanreturn output_image[0].clamp(0, 1)if __name__ == '__main__':train()

重点可以看看这段代码，按照代码的解释，这里上采样不是通过convTranspose来实现，而是先做一个ReflectionPad，再进行插值来上采样，然后做一个conv，按照注释的说法，这样效果比ConvTranspose2d好。

class UpsampleConvLayer(torch.nn.Module):"""UpsampleConvLayerUpsamples the input and then does a convolution. This method gives better resultscompared to ConvTranspose2d.ref: http://distill.pub/2016/deconv-checkerboard/"""def __init__(self, in_channels, out_channels, kernel_size, stride, upsample=None):super(UpsampleConvLayer, self).__init__()self.upsample = upsamplereflection_padding = kernel_size // 2self.reflection_pad = torch.nn.ReflectionPad2d(reflection_padding) # 可以重点关注这个padding方式,参考https://blog.csdn.net/LionZYT/article/details/120181586self.conv2d = torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride)def forward(self, x):x_in = xif self.upsample:x_in = torch.nn.functional.interpolate(x_in, mode='nearest', scale_factor=self.upsample)out = self.reflection_pad(x_in)out = self.conv2d(out)return out