【PyTorch】PyTorch之风格迁移

前面介绍了PyTorch在深度学习上的一些简单应用，这一节讲解PyTorch在风格迁移上的用法。

基础知识

numpy.array()

• 将矩阵或者拥有$__array__$方法的对象或者sequence转化为矩阵。

array.astype()

• 将矩阵转化为对应的数据类型。

Tensor.squeeze()

• 若不指定dim，则将Tensor中dim=1的维度给移除；若指定dim，则只将指定的dim=1的维度给移除。

Tensor.unsqueeze()

• 在指定的维度上插入dim=1。

Tensor.type()

• 如果不带参数，则返回该Tensor的数据类型；否则将该Tensor转化为指定的数据类型。

Tensor.mean()

• 如果指定维度，则计算该维度上的均值，返回Tensor；否则计算全局均值，返回float。

Tensor.mm()

• 计算两个Tensor的矩阵乘法，返回Tensor。

Tensor.clamp()

• 将Tensor中的所有元素clamp到[min, max]。

风格迁移实战

其实要实现的东西很清晰，就是需要将两张图片融合在一起，这个时候就需要定义怎么才算融合在一起。首先需要的就是内容上是相近的，然后风格上是相似的。这样来我们就知道我们需要做的事情是什么了，我们需要计算融合图片和内容图片的相似度，或者说差异性，然后尽可能降低这个差异性；同时我们也需要计算融合图片和风格图片在风格上的差异性，然后也降低这个差异性就可以了。这样我们就能够量化我们的目标了。

对于内容的差异性我们该如何定义呢？其实我们能够很简答的想到就是两张图片每个像素点进行比较，也就是求一下差，因为简单的计算他们之间的差会有正负，所以我们可以加一个平方，使得差全部是正的，也可以加绝对值，但是数学上绝对值会破坏函数的可微性，所以大家都用平方，这个地方不理解也没关系，记住普遍都是使用平方就行了。

对于风格的差异性我们该如何定义呢？这才是一个难点。这也是这篇文章提出的创新点，引入了Gram矩阵计算风格的差异。我尽量不使用数学的语言来解释，而使用通俗的语言。

Gram矩阵是如何定义的呢？首先Gram矩阵的大小是有特征图的厚度决定的，等于 CxC，那么每一个Gram矩阵的元素，也就是 Gram(i, j) 等于多少呢？先把特征图中第 i 层和第 j 层取出来，这样就得到了两个 MxN的矩阵，然后将这两个矩阵对应元素相乘然后求和就得到了 Gram(i, j)，同理 Gram 的所有元素都可以通过这个方式得到。这样 Gram 中每个元素都可以表示两层特征图的一种组合，就可以定义为它的风格。

作者：SherlockLiao 链接：https://www.jianshu.com/p/8f8fc2aa80b3 來源：简书
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

import torch
import torch.nn as nn
from torch.autograd import Variable
import torchvision
from torchvision import transforms, models
from PIL import Image
import argparse
import numpy as np
import os

全局变量，是否用gpu。

use_gpu = torch.cuda.is_available()
dtype = torch.cuda.FloatTensor if use_gpu else torch.FloatTensor

定义加载图像函数，并将PIL image转化为Tensor。

def load_image(image_path, transforms=None, max_size=None, shape=None):
    image = Image.open(image_path)
    image_size = image.size

    if max_size is not None:
        #获取图像size，为sequence
        image_size = image.size
        #转化为float的array
        size = np.array(image_size).astype(float)
        size = max_size / size * size;
        image = image.resize(size.astype(int), Image.ANTIALIAS)

    if shape is not None:
        image = image.resize(shape, Image.LANCZOS)

    #必须提供transform.ToTensor，转化为4D Tensor
    if transforms is not None:
        image = transforms(image).unsqueeze(0)

    #是否拷贝到GPU
    return image.type(dtype)

定义VGG19模型，前向时提取第0，5，10，19, 28层卷积特征。

class VGGNet(nn.Module):
    def __init__(self):
        super(VGGNet, self).__init__()
        self.select = ['0', '5', '10', '19', '28']
        self.vgg19 = models.vgg19(pretrained = True).features

    def forward(self, x):
        features = []
        #name类型为str，x为Variable
        for name, layer in self.vgg19._modules.items():
            x = layer(x)
            if name in self.select:
                features.append(x)
        return features

定义主函数，获取5个卷积层对应的content和style特征，分别计算content_loss和style_loss。

def main(config):
    #定义图像变换操作，必须定义.ToTensor()。
    transform = transforms.Compose(
        [transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), 
                             (0.229, 0.224, 0.225))
        ])

    #加速content和style图像，style图像resize成同样大小
    content = load_image(config.content, transform, max_size = config.max_size)
    style = load_image(config.style, transform, shape = [content.size(2), content.size(3)])

    #将concent复制一份作为target，并需要计算梯度，作为最终的输出
    target = Variable(content.clone(), requires_grad = True)
    optimizer = torch.optim.Adam([target], lr = config.lr, betas=[0.5, 0.999])

    vgg = VGGNet()
    if use_gpu:
        vgg = vgg.cuda()

    for step in range(config.total_step):
        #分别计算5个特征图
        target_features = vgg(target)
        content_features = vgg(Variable(content))
        style_features = vgg(Variable(style))

        content_loss = 0.0
        style_loss = 0.0

        for f1, f2, f3 in zip(target_features, content_features, style_features):
            #计算content_loss
            content_loss += torch.mean((f1 - f2)**2)

            n, c, h, w = f1.size()

            #将特征reshape成二维矩阵相乘，求gram矩阵
            f1 = f1.view(c, h * w)
            f3 = f3.view(c, h * w)

            f1 = torch.mm(f1, f1.t())
            f3 = torch.mm(f3, f3.t())

            #计算style_loss
            style_loss += torch.mean((f1 - f3)**2) / (c * h * w) 

        #计算总的loss
        loss = content_loss + style_loss * config.style_weight

        #反向求导与优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (step+1) % config.log_step == 0:
            print ('Step [%d/%d], Content Loss: %.4f, Style Loss: %.4f' 
                   %(step+1, config.total_step, content_loss.data[0], style_loss.data[0]))

        if (step+1) % config.sample_step == 0:
            # Save the generated image
            denorm = transforms.Normalize((-2.12, -2.04, -1.80), (4.37, 4.46, 4.44))
            img = target.clone().cpu().squeeze()
            img = denorm(img.data).clamp_(0, 1)
            torchvision.utils.save_image(img, 'output-%d.png' %(step+1))

从命令行获取参数。

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--content', type=str, default='/home/content.jpg')
    parser.add_argument('--style', type=str, default='/home/style.jpg')
    parser.add_argument('--max_size', type=int, default=400)
    parser.add_argument('--total_step', type=int, default=5000)
    parser.add_argument('--log_step', type=int, default=10)
    parser.add_argument('--sample_step', type=int, default=1000)
    parser.add_argument('--style_weight', type=float, default=100)
    parser.add_argument('--lr', type=float, default=0.003)
    config = parser.parse_args()
    print(config)
    main(config)

实验结果为：