本文给出简单代码实现风格迁移。
1，原理简介
风格迁移和上篇文章提到的deep dream算法比较接近，都是根据某种优化指标计算梯度来反向优化输入图像的像素。所以在学完deep dream之后趁热打铁又学了这个，但本文仅限于基础版的实现，对该领域后来发展出的诸多进化版不做讨论。
基于深度学习的风格迁移最早由 Gatys于2015年提出，其核心理论是使用格拉姆矩阵（gram matrix）来提取图像的风格。gram matrix的计算方法是对图像的某中间层特征features（尺寸为CHW），reshape为C*(HW)后，再和自己的转置做矩阵乘法，就可以得到一个CC的矩阵，这个矩阵是原矩阵的偏心协方差矩阵，它去掉了像素级的信息，而表达了各通道之间的相关性。我们从deep dream的分析中已经认识到，中间层特征的各通道表达了图像特征的不同维度，比如有的表示尖顶建筑，有的表示黑色条纹，等等。如果我们把两个通道特征按像素顺序相乘再求总和，就表示这两种特征同时出现或同时不出现的程度。比如尖顶建筑总是和黑色条纹同时出现，那么这就是图像的某种偏好，也就是它的风格。从这个分析中我们可以看出，这样定义的风格是与图像中特征出现的位置无关的，但与两种特征是否同时出现有关，这个定义非常巧妙。但我总觉得这个定义也并不是非常完备，这应该只是一种简单的风格，而更多实际风格应该比这要复杂的多。
不管怎么说，从数量上提取出了“风格”这个量之后，我们就可以操作它了，可以玩出各种花，最常用的是我们把风格从一张图像迁移到另一张图像上。算法原理如下图（二手中转自知乎）：
先选择一张风格图片 style image，和一张内容图片 content image，然后还需要给定一张初始种子图片，可以是一张噪声图，也可以就是content image，当然也可以是一张其他图片，效果各不相同。然后先让风格图片和内容图片从网络中过一遍，风格图片则从网络的多个中间层提取出特征图，计算各个层的风格矩阵（gram matrix），因为浅层的风格比较基础，深层的风格比较高级，都有用。（实际上我觉得把浅层和深层也交叉计算一下gram matrix应该也有用，因为有时候浅层和深层特征也会同时出现，这个留作有空做一下试验）。然后再把内容图片从网络中过一下，只提取其中后层的某层特征，因为浅层特征本身就含有更多的风格意味，我们迁移风格的时候不需要保留它们。最后我们再把初始图片从网络中迭代的循环通过，每次正向通过时计算各对应层的风格，和风格图的风格比较得到风格损失，计算对应层的内容，和内容图比较得到内容损失。论文中作者又加入了第三种损失，即下面代码中的TV loss，表示生成图像的光滑程度，防止图像失真。根据综合损失求输入图的梯度，即可以迭代的优化输入图，得到一张风格和风格图一致，内容和内容图一致，又保有种子图片结构（如果种子图不是随机噪声）的新的图片。
下面还是看完整代码吧。
2，完整代码
import torch
import torchvision.models as models
import torch.nn.functional as F
import torch.nn as nn
import numpy as np
import numbers
import math
import cv2
from PIL import Image
from torchvision.transforms import Compose, ToTensor, Normalize, Resize, ToPILImage
import time
t0 = time.time()
model = models.vgg19(pretrained=True).cuda()
batch_size = 1
for params in model.parameters():
params.requires_grad = False
model.eval()
mu = torch.Tensor([0.485, 0.456, 0.406]).unsqueeze(-1).unsqueeze(-1).cuda()
std = torch.Tensor([0.229, 0.224, 0.225]).unsqueeze(-1).unsqueeze(-1).cuda()
unnormalize = lambda x: x*std + mu
normalize = lambda x: (x-mu)/std
transform_test = Compose([
Resize((512,512)),
ToTensor(),
])
content_img = Image.open('./data/tubingen.jpg')
image_size = content_img.size
content_img = transform_test(content_img).unsqueeze(0).cuda()
style_img = Image.open('./data/starry_night.jpg')
style_img = transform_test(style_img).unsqueeze(0).cuda()
var_img = content_img.clone()
#var_img = torch.rand_like(content_img)
var_img.requires_grad=True
class ShuntModel(nn.Module):
def __init__(self, model):
super().__init__()
self.module = model.features.cuda().eval()
self.con_layers = [22]
self.sty_layers = [1,6,11,20,29]
for name, layer in self.module.named_children():
if isinstance(layer, nn.MaxPool2d):
self.module[int(name)] = nn.AvgPool2d(kernel_size = 2, stride = 2)
def forward(self, tensor: torch.Tensor) -> dict:
sty_feat_maps = []; con_feat_maps = [];
x = normalize(tensor)
for name, layer in self.module.named_children():
x = layer(x);
if int(name) in self.con_layers: con_feat_maps.append(x)
if int(name) in self.sty_layers: sty_feat_maps.append(x)
return {"Con_features": con_feat_maps, "Sty_features": sty_feat_maps}
model = ShuntModel(model)
sty_target = model(style_img)["Sty_features"]
con_target = model(content_img)["Con_features"]
gram_target = []
for i in range(len(sty_target)):
b, c, h, w = sty_target[i].size()
tensor_ = sty_target[i].view(b * c, h * w)
gram_i = torch.mm(tensor_, tensor_.t()).div(b*c*h*w)
gram_target.append(gram_i)
optimizer = torch.optim.Adam([var_img], lr = 0.01, betas = (0.9,0.999), eps = 1e-8)
lam1 = 1e-3; lam2 = 1e7; lam3 = 5e-3
for itera in range(20001):
optimizer.zero_grad()
output = model(var_img)
sty_output = output["Sty_features"]
con_output = output["Con_features"]
con_loss = torch.tensor([0]).cuda().float()
for i in range(len(con_output)):
con_loss = con_loss + F.mse_loss(con_output[i], con_target[i])
sty_loss = torch.tensor([0]).cuda().float()
for i in range(len(sty_output)):
b, c, h, w = sty_output[i].size()
tensor_ = sty_output[i].view(b * c, h * w)
gram_i = torch.mm(tensor_, tensor_.t()).div(b*c*h*w)
sty_loss = sty_loss + F.mse_loss(gram_i, gram_target[i])
b, c, h, w = style_img.size()
TV_loss = (torch.sum(torch.abs(style_img[:, :, :, :-1] - style_img[:, :, :, 1:])) +
torch.sum(torch.abs(style_img[:, :, :-1, :] - style_img[:, :, 1:, :])))/(b*c*h*w)
loss = con_loss * lam1 + sty_loss * lam2 + TV_loss * lam3
loss.backward()
var_img.data.clamp_(0, 1)
optimizer.step()
if itera%100==0:
print('itera: %d, con_loss: %.4f, sty_loss: %.4f, TV_loss: %.4f'%(itera,
con_loss.item()*lam1,sty_loss.item()*lam2,TV_loss.item()*lam3),'\n\t total loss:',loss.item())
print('var_img mean:%.4f, std:%.4f'%(var_img.mean().item(),var_img.std().item()))
print('time: %.2f seconds'%(time.time()-t0))
if itera%1000==0:
save_img = var_img.clone()
save_img = torch.clamp(save_img,0,1)
save_img = save_img[0].permute(1,2,0).data.cpu().numpy()*255
save_img = save_img[...,::-1].astype('uint8') #注意cv2使用BGR顺序
save_img = cv2.resize(save_img,image_size)
cv2.imwrite('./data/output1/transfer%d.jpg'%itera,save_img)
3，效果
该方法的速度较慢，通常要好几分钟才能生成一幅图片。后续的改进论文中有快速生成算法，也有按照区域迁移（例如避免把天空中的风格迁移到建筑上）使效果更好等等诸多改进，暂时先不打算学了。（这两年AI坑挖的太快，知识量太多，只能先粗学一遍）
这个技术还多少有点用处，即使用本文这个基础版代码，在网上随便找一些风格图片简单处理一下，也可以生成一些还不错的图片，下面是我宝贝女儿的照片进行处理后的一些结果。
4，补充（2021.4.13）
上面提到可以把多层的特征合并后统一计算格拉姆矩阵，这个实现并不难，所以今天顺手就试了一下，发现有效。
和图3最右图对比，可以发现，使用统一格拉姆矩阵时转移了风格图上更多更细节的风格，比如：鼻尖上出现黄点，眼睛、眉毛的形状更像风格图，脸上涂色更像风格图，等。但是这种方法更加占用显存（内存），计算速度也慢很多，根据具体需要选择吧。