【学习笔记】pytorch-图像风格迁移实战

1. 风格迁移入门
我们首先选取一幅图像作为基准图像，也可以将其称为内容图像，然后选取另一幅或多幅图像作为我们希望获取相应风格的图像，也可以将其称为风格图像。图像风格迁移算法就是在保证内容图像的完整性的前提下，将风格图像的风格融入内容图像中，使得内容图像的原始风格最后发生了转变，最终输出的图像呈现的将是输入的内容图像的内容和风格图像风格之间的理想融合。
图像风格迁移实现的难点就是如何有效地提取一张图像的风格。和传统的图像风格提取方法不同，我们在基于神经网络的图像风格迁移算法中使用卷积神经网络来完成对图像风格的提取。

2. pytorch实战
首先我们需要获取一张内容图片和一张风格图片；然后定义两个度量值，一个度量叫做内容度量值，另一个叫做风格度量值，其中的内容度量值用于衡量图片之间的内容差异程度，风格度量值用于衡量图片之间的风格差异程度；最后，简历神经网络模型，对内容图片中的内容和风格图片的风格进行提取，以内容图片为基准将其输入建立的模型中，并不断调整内容度量值和风格度量值，让他们趋近最小，最后输出的图片就是内容与风格融合的图片。

2.1 图像的内容损失
内容度量值可以使用均方误差作为损失函数，在代码中定义的图像内容损失如下：

class Content_loss(torch.nn.Module):
	def __init__(self,weight,targrt):
		super(Content_loss, self).__init__()
		self.weight = weight
		self.target = target.detach()*weight
		self.loss_fn = torch.nn.MSELoss()
		
	def forward(self, input):
		self.loss = self.loss_fn(input*self.weight, self.target)
		return input
	def backward(self):
		self.loss.backward(retain_graph = True)
		return self.loss

以上代码中的target是通过卷积获取到的输入图像中的内容，weight是我们设置的一个权重参数，用来控制内容和风格对最后合成图像的影响程度；input代表输入图像，target.detach()用于对提取到的内容进行锁定，不需要进行梯度；forward函数用于计算输入图像和内容图像之间的损失值；backward函数根据计算得到的损失值进行后向传播，并返回损失值。

2.2 图像的风格损失
风格度量同样使用均方误差作为损失函数，代码如下：

class Style_loss(torch.nn.Module):
	def __init__(self,weight,target):
		super(Style_loss, self).__init__():
		self.weight = weight
		self.target = target
		self.loss_fn = torch.nn.MSELoss()
		self.gram = Gram_matrix()
	def forward(self,input):
		self.Gram = self.gram(input.clone())
		self.Gram.mul_(self.weight)
		self.loss = self.loss_fn(self.Gram, self.target)
		return input
	def backward(self):
		self.loss.backward(retain_graph = True)
		return self.loss

风格损失计算的代码基本和内容损失计算的代码相似，不同之处是在代码中引入了一个Gram_matri类定义的实例参与风格损失的计算，这个类的代码如下：

class Gram_matrix(torch.nn.Module):
	def forward(self,input):
		a,b,c,d = input.size()
		feature = input.view(a*b,c*d)
		gram = torch.mm(feature, feature.t())
		return gram.div(a*b*c*d)

以上代码实现的是格拉姆矩阵的功能。我们通过卷积神经网络提取了风格图片的风格，这些风格其实是由数字组成的，数字的大小代表了图中风格的突出程度，而Gram矩阵是矩阵的内积运算，在运算过后输入到该矩阵的特征图中的大的数字会变得更大，相当于图片的风格被放大了，放大的风格在参与损失计算，便能够对最后的合成图片产生更大的影响。

2.3 模型搭建和参数优化
在定义好内容损失和风格损失的计算方法之后，我们还需要搭建一个自定义的模型，并将这两部分内容融入模型中。我们首先要做的是欠以一个卷积神经网络的特征提取部分，即卷积相关的部分，代码如下：

cnn = models.vgg16(pretrained = True).features

content_layer = ["Conv_3"]
style_layer = ["Conv_1","Conv_2","Conv_3","Conv_4"]

content_losses = []
style_losses = []

conten_weight = 1
style_weight = 1000

以上代码迁移了一个vgg16架构的卷积神经网络模型的特征提取部分，然后定义了content_layer和style_layer，分别制定了我们需要在整个卷积过程中的哪一层提取内容，以及在哪一层提取风格。content_loss和style_loss是两个用于保存内容损失和风格损失的列表；content_weight和style_weight指定了内容损失和风格损失对我们最后得到的融合图片的影响权重。
搭建图像风格迁移模型的代码如下：

new_model = torch.nn.Sequential()
model = copy.deepcopy(cnn)
gram = gram_matrix()

if use_gpu:
	new_model = new_model.cuda()
	gram = gram.cuda()
index = 1
for layer in list(model)[:8]:
	if isinstance(layer, torch.nn.Conv2d):
		name = "Conv_"+str(index)
		new_model.add_module(name,layer)
		for name in content_layer:
			target = new_model(content_img).clone()
			content_loss = Content_loss(content_weight, target)
			new_model.add_module("content_loss_"+str(index),content_loss)
			content_losses.append(content_loss)

		if name in style_layer:
			target = new_model(style_img).clone()
			target = gram(target)
			style_loss = Style_loss(style_weight, target)
			new_model.add_module("style_loss_"+str(index), style_loss)
			style_losses.append(style_loss)
	
	if isinstance(layer, torch.nn.ReLU):
		name = "ReLU_"+str(index)
		new_model.add_module(name, layer)
		index = index + 1
	if isinstance(layer, torch.nn.MaxPool2d):
		name = "MaxPool_"+str(index)
		new_model.add_module(name, layer)

以上代码中，for layer in list(model)[:8]指明了我们仅仅用到迁移模型特征提取部分的前8层，，因为我们的内容提取和风格提取在前8层就已经完成了。然后建立一个空的模型，使用torch.nn.Module类的add_module方法向空的模型中加入指定的层次模块，最后得到我们自定义的图像风格迁移模型。add_module方法传递的参数分别是层次的名字和模块，该模块是使用isinstance实例检测函数得到的，而名字是对应的层次。在定义好模型之后可以对其进行打印输出。
接下来是参数优化部分的代码。

input_image = content_img.clone()
parameter = torch.nn.Parameter(input_img.data)
optimizer = torch.optim.LBFGS([parameter])

在以上代码中使用的优化函数是optim.LBFGS，原因是这个模型中需要优化的损失值有多个并且规模较大，使用该优化函数可以取得更好的效果。

2.4 训练新定义的卷积神经网络
完成模型的搭建和优化函数定义之后，就可以开始进行模型的训练和参数的优化，代码如下：

epoch_n = 300

epoch = [0]
while epoch[0]<epoch_n:
	def closure():
		optimizer.zero_grad()
		style_score = 0
		content_score = 0
		parameter.data.clamp_(0,1)
		new_model(parameter)
		for sl in style_losses:
			style_score += sl.backward()
		
		for cl in content_losses:
			content_score += cl.backward()
		
		epoch[0]+=1
		if epoch[0]%50 ==0:
			print("Epoch:{} Style score:{:.4f} Content Score:{:.4f}".format(epoch[0],style_score.data,content_score.data))
		return style_score+content_score
	optimizer.step(closure)

2.5 完整图像风格迁移代码

import torch
import torchvision
from torchvision import transforms, models
from PIL import Image
import matplotlib.pyplotlib.pyplot as plt
from torch.autograd import Variable
import copy

%matplotlib inline

transform = transforms.Compose([transforms.Scale([224,224]),transforms.ToTensor()])

def loading(path = None):
	img = Image.open(path)
	img = transform(img)
	img = img.unsqueeze(0)
	return img

content_img = loading("images/4.jpg")
content_img = Variable(content_img).cuda()
style_img = loading("images/1.jpg")
style_img = Variable(style_img).cuda()

class Content_loss(torch.nn.Module):
	def __init__(self,weight,targrt):
		super(Content_loss, self).__init__()
		self.weight = weight
		self.target = target.detach()*weight
		self.loss_fn = torch.nn.MSELoss()
		
	def forward(self, input):
		self.loss = self.loss_fn(input*self.weight, self.target)
		return input
	def backward(self):
		self.loss.backward(retain_graph = True)
		return self.loss
class Gram_matrix(torch.nn.Module):
	def forward(self,input):
		a,b,c,d = input.size()
		feature = input.view(a*b,c*d)
		gram = torch.mm(feature, feature.t())
		return gram.div(a*b*c*d)
class Style_loss(torch.nn.Module):
	def __init__(self,weight,target):
		super(Style_loss, self).__init__():
		self.weight = weight
		self.target = target
		self.loss_fn = torch.nn.MSELoss()
		self.gram = Gram_matrix()
	def forward(self,input):
		self.Gram = self.gram(input.clone())
		self.Gram.mul_(self.weight)
		self.loss = self.loss_fn(self.Gram, self.target)
		return input
	def backward(self):
		self.loss.backward(retain_graph = True)
		return self.loss
use_gpu = torch.cuda.is_available()
cnn = models.vgg16(pretrained = True).features

if use_gpu:
	cnn = cnn.cuda()
model = copy.deepcopy(cnn)

content_layer = ["Conv_3"]
style_layer = ["Conv_1","Conv_2","Conv_3","Conv_4"]

content_losses = []
style_losses = []

conten_weight = 1
style_weight = 1000

new_model = torch.nn.Sequential()
model = copy.deepcopy(cnn)
gram = gram_matrix()

if use_gpu:
	new_model = new_model.cuda()
	gram = gram.cuda()
index = 1
for layer in list(model)[:8]:
	if isinstance(layer, torch.nn.Conv2d):
		name = "Conv_"+str(index)
		new_model.add_module(name,layer)
		for name in content_layer:
			target = new_model(content_img).clone()
			content_loss = Content_loss(content_weight, target)
			new_model.add_module("content_loss_"+str(index),content_loss)
			content_losses.append(content_loss)

		if name in style_layer:
			target = new_model(style_img).clone()
			target = gram(target)
			style_loss = Style_loss(style_weight, target)
			new_model.add_module("style_loss_"+str(index), style_loss)
			style_losses.append(style_loss)
	
	if isinstance(layer, torch.nn.ReLU):
		name = "ReLU_"+str(index)
		new_model.add_module(name, layer)
		index = index + 1
	if isinstance(layer, torch.nn.MaxPool2d):
		name = "MaxPool_"+str(index)
		new_model.add_module(name, layer)

input_image = content_img.clone()
parameter = torch.nn.Parameter(input_img.data)
optimizer = torch.optim.LBFGS([parameter])

epoch_n = 300

epoch = [0]
while epoch[0]<epoch_n:
	def closure():
		optimizer.zero_grad()
		style_score = 0
		content_score = 0
		parameter.data.clamp_(0,1)
		new_model(parameter)
		for sl in style_losses:
			style_score += sl.backward()
		
		for cl in content_losses:
			content_score += cl.backward()
		
		epoch[0]+=1
		if epoch[0]%50 ==0:
			print("Epoch:{} Style score:{:.4f} Content Score:{:.4f}".format(epoch[0],style_score.data,content_score.data))
		return style_score+content_score
	optimizer.step(closure)

3. 小结
本章展示的是比较基础的图像风格迁移算法，每次训练只能对其中的一种风格进行迁移，如果要进行其他风格的迁移，则需要调整模型重新训练。