用 PyTorch 做 神经转换 (Neural Transfer)

译者：@Twinkle

原作者: Alexis Jacq

介绍

欢迎观看! 这篇教程解释了如何实现 Leon A. Gatys, Alexander S. Ecker 和 Matthias Bethge 几位学者发明的 Neural-Style 算法 .

题中的神经描述的是什么?

神经风格, 或者说神经转换是一种算法, 它输入一张内容图像 (例如海龟), 一张风格图像 (例如艺术波浪), 然后返回内容图像的内容, 此时返回的内容像是被艺术风格图像的风格渲染过:

content1

它是如何工作的?

原理很简单: 我们定义两个距离, 一个是关于内容的 (D_C) , 另一个是关于风格的 (D_S) . D_C 衡量两张图像的内容有多么不同, 而 D_S 衡量两张图像的风格有多么不同. 接着我们拿出我们的输入, 也就是第三张图像 (例如全噪声), 然后我们转换它, 同时最小化它与内容图像的内容距离和它与风格图像的风格距离.

好吧, 它具体是怎么工作的?

继续深入需要一些数学知识. 令 C_{nn} 代表一个预训练好的深度卷积神经网络, X 代表任何图像. C_{nn}(X) 是神经网络输入 X 后的结果 (包括在所有层的特征映射). 令 F_{XL} \in C_{nn}(X) 代表在深度为 L 层处的特征映射, 都矢量化和级联为一个单一矢量. 我们简单地用 F_{XL} 定义 X 在 L 层的内容. 如果 Y 是另一张和 X 相同大小的图像, 我们定义这两张图像在 L 层的内容距离如下:

D_C^L(X,Y) = \|F_{XL} - F_{YL}\|^2 = \sum_i (F_{XL}(i) - F_{YL}(i))^2

式中 F_{XL}(i) 是 F_{XL} 的第 i^{th} 个元素. 定义风格要更繁琐一些. 令满足 k \leq K 的 F_{XL}^k 代表 L 层矢量化的 K 个特征映射中的第 k^{th} 个. 图像 X 在 L 层的风格 G_{XL} 定义为满足 k \leq K 的所有矢量化特征映射 F_{XL}^k 的克产物 (Gram produce). 换句话说, G_{XL} 是一个 KxK 的矩阵, 其在 k^{th} 行和 l^{th} 列的每个元素 G_{XL}(k,l) 是 F_{XL}^k 和 F_{XL}^l 之间的矢量产物 :

G_{XL}(k,l) = \langle F_{XL}^k, F_{XL}^l\rangle = \sum_i F_{XL}^k(i) . F_{XL}^l(i)

式中 F_{XL}^k(i) 是 F_{XL}^k 的第 i^{th} 个元素. 我们可以把 G_{XL}(k,l) 当做特征映射 k 和 l 相关性的衡量. 那样的话, G_{XL} 代表了 X 在 L 层特征向量的相关性矩阵. 注意 G_{XL} 的尺寸只决定于特征映射的数量, 不被 X 的尺寸所影响. 然后如果 Y 是 任意尺寸 的另一张图像, 我们定义在 L 层的风格距离如下:

D_S^L(X,Y) = \|G_{XL} - G_{YL}\|^2 = \sum_{k,l} (G_{XL}(k,l) - G_{YL}(k,l))^2

要想一次性地在一些层最小化一个可变图像 X 与目标内容图像 C 间的 D_C(X,C), 和 X 与目标风格图像 S 间的 D_S(X,S) , 我们计算并加和每个目标层每个距离的梯度 (对 X 求导).

\nabla_{ extit{total}}(X,S,C) = \sum_{L_C} w_{CL_C}.\nabla_{ extit{content}}^{L_C}(X,C) + \sum_{L_S} w_{SL_S}.\nabla_{ extit{style}}^{L_S}(X,S)

式中 L_C 和 L_S 分别是内容和风格的目标层(任意陈述), w_{CL_C} 和 w_{SL_S} 是风格和内容关于每个目标层的权重(任意陈述). 然后我们对 X 进行梯度下降:

X \leftarrow X - \alpha \nabla_{ extit{total}}(X,S,C)

好吧, 数学的部分就到此为止. 如果你想要更加深入(比如怎么计算梯度), 我们推荐你阅读原始论文 (作者是 Leon A. Gatys 和 AL), 论文中这部分解释地更好更清晰.

对于在 PyTorch 中的实现, 我们已经有了我们需要的一切: 事实上就是 PyTorch, 所有的梯度都被为你自动且动态地计算(当你从库中使用函数时). 这就是为什么算法的实现在 PyTorch 中变得非常轻松.

PyTorch 实现

如果你不确定是否理解了以上数学公式, 你也可以通过实现它, 在过程中有所领悟. 如果你正在探索 PyTorch , 我们推荐你先阅读这篇教程 Introduction to PyTorch.

包

我们将会依赖下列这些包:

  • torch, torch.nn, numpy (indispensables packages for neural networks with PyTorch)
  • torch.autograd.Variable (dynamic computation of the gradient wrt a variable)
  • torch.optim (efficient gradient descents)
  • PIL, PIL.Image, matplotlib.pyplot (load and display images)
  • torchvision.transforms (treat PIL images and transform into torch tensors)
  • torchvision.models (train or load pre-trained models)
  • copy (to deep copy the models; system package)
  1. from __future__ import print_function
  3. import torch
  4. import torch.nn as nn
  5. from torch.autograd import Variable
  6. import torch.optim as optim
  8. from PIL import Image
  9. import matplotlib.pyplot as plt
  11. import torchvision.transforms as transforms
  12. import torchvision.models as models
  14. import copy

Cuda

如果你的计算机里有 GPU, 推荐在上面运行算法, 尤其是当你要尝试 大型网络时 (就像 VGG). 有鉴于此, 我们有 torch.cuda.is_available(), 如果你的计算机有可用 GPU 则会返回 True. 然后我们用 .cuda() 方法 将可分配的进程和模块从 CPU 移动到 GPU. 当我们想将这些模块重新移回 CPU 的时候(比如要用 numpy), 我们用 .cpu() 方法. 最后, .type(dtype) 会用来将一个 torch.FloatTensor 转化为 用于 GPU 进程输入的 torch.cuda.FloatTensor.

  1. use_cuda = torch.cuda.is_available()
  2. dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor

读取图像

为了简化其实现, 让我们从导入一个相同维度的风格和内容图像开始. 然后我们将它们缩放到想要的输入图像尺寸 (在例子中是 128 和 512, 取决你的 GPU 是否可用) 然后把它们转化为 torch 张量, 以待喂入一个神经网络.

注解：

这里是教程需要的图像的下载链接: picasso.jpg 和 dancing.jpg. 下载这两张图像然后把它们加入到名为 images 的目录中.

  1. # 想要的输出图像尺寸
  2. imsize = 512 if use_cuda else 128  # 如果没有 GPU 则使用小尺寸
  4. loader = transforms.Compose([
  5.     transforms.Scale(imsize),  # 缩放图像
  6.     transforms.ToTensor()])  # 将其转化为 torch 张量
  8. def image_loader(image_name):
  9.     image = Image.open(image_name)
  10.     image = Variable(loader(image))
  11.     # 由于神经网络输入的需要, 添加 batch 的维度
  12.     image = image.unsqueeze(0)
  13.     return image
  15. style_img = image_loader("images/picasso.jpg").type(dtype)
  16. content_img = image_loader("images/dancing.jpg").type(dtype)
  18. assert style_img.size() == content_img.size(), \
  19.     "we need to import style and content images of the same size"

导入的 PIL 图像像素值的范围为 0 到 255. 转化为 torch 张量后, 它们的值范围变为了 0 到and 1. 这是个重要的细节: torch 库中的神经网络被使用 0-1 的张量图像训练. 如果你尝试用 0-255 的张量图像喂入神经网络, 激活的特征映射就没用了. 这不是 使用 Caffe 库中预训练的神经网络, Caffe 中是用 0-255 的张量图像训练的.

显示图像

我们将使用 plt.imshow 来显示图像. 所以我们需要先把它们转回 PIL 图像.

  1. unloader = transforms.ToPILImage()  # 转回 PIL 图像
  3. plt.ion()
  5. def imshow(tensor, title=None):
  6.     image = tensor.clone().cpu()  # 克隆是为了不改变它
  7.     image = image.view(3, imsize, imsize)  # 移除 batch 维度
  8.     image = unloader(image)
  9.     plt.imshow(image)
  10.     if title is not None:
  11.         plt.title(title)
  12.     plt.pause(0.001) # 暂停一会, 让绘图更新
  14. plt.figure()
  15. imshow(style_img.data, title='Style Image')
  17. plt.figure()
  18. imshow(content_img.data, title='Content Image')

内容损失

内容损失是一个在网络输入为 X 的层 L 输入特征映射 F_{XL} 的函数, 返回此图像与内容图像间的加权内容距离 w_{CL}.D_C^L(X,C). 之后, 权重 w_{CL} 和目标内容 F_{CL} 就成为了函数的参数. 我们把这个函数作为 torch 模块来实现, 把这些参 数作为构造器的输入. 这个距离 \|F_{XL} - F_{YL}\|^2 是两个特征映射集的 均方误差, 可以用作为第三个参数的标准的 nn.MSELoss 来计算.

我们会在每个目标层加入我们的内容损失作为额外的神经网络模块. 这样, 每次我们都会给神经 网络投喂一张输入图像 X, 所有的损失都会在目标层被计算, 多亏了自动梯度计算, 所有梯度都会被搞定. 要实现, 我们只需写出转换模块的 forward 方法, 这个模块就变 成了网络的 ‘’transparent layer (透明层)’‘, 计算好的损失被存为模块的参数.

最后, 我们定义一个假的 backward 方法, 它仅仅只调用后向方法 nn.MSELoss 来重构梯度. 这个方法返回计算好的损失: 运行梯度下降时要想显示风格和内容损失的变化, 这会非常有用.

  1. class ContentLoss(nn.Module):
  3.     def __init__(self, target, weight):
  4.         super(ContentLoss, self).__init__()
  5.         # 我们会从所使用的树中“分离”目标内容
  6.         self.target = target.detach() * weight
  7.         # 动态地计算梯度: 它是个状态值, 不是变量.
  8.         # 否则评价指标的前向方法会抛出错误.
  9.         self.weight = weight
  10.         self.criterion = nn.MSELoss()
  12.     def forward(self, input):
  13.         self.loss = self.criterion(input * self.weight, self.target)
  14.         self.output = input
  15.         return self.output
  17.     def backward(self, retain_graph=True):
  18.         self.loss.backward(retain_graph=retain_graph)
  19.         return self.loss

注解：

重要细节: 这个模块虽然叫做 ContentLoss, 却不是个真正的 Pytorch 损失函数. 如果你想像 Pytorch 损失一样定义你的内容损失, 你得新建一个 Pytorch 自动求导函数并手动得在 backward 方法中重算/实现梯度.

风格损失

对于风格损失, 我们首先需要定义一个给定输入 X 在 L 层的特征映射 F_{XL} 时计算克产物 G_{XL} 的模块. 令 \hat{F}_{XL} 表示 F_{XL} 重变形为 KxN 的版本, 这里 K 是 L 层特征 映射的数量, N 是任意矢量化特征映射 F_{XL}^k 的长度. \hat{F}_{XL} 的第 k^{th} 行是 F_{XL}^k. 可以验证 \hat{F}_{XL} \cdot \hat{F}_{XL}^T = G_{XL}. 鉴于此, 实现我们的模块就很容易了:

  1. class GramMatrix(nn.Module):
  3.     def forward(self, input):
  4.         a, b, c, d = input.size()  # a=batch size(=1)
  5.         # b= 特征映射的数量
  6.         # (c,d)= 一个特征映射的维度 (N=c*d)
  8.         features = input.view(a * b, c * d)  # 将 F_XL 转换为 \hat F_XL
  10.         G = torch.mm(features, features.t())  # 计算克产物 (gram product)
  12.         # 我们用除以每个特征映射元素数量的方法
  13.             # 标准化克矩阵 (gram matrix) 的值
  14.         return G.div(a * b * c * d)

特征映射的维度 N 越长, 则克矩阵 (gram matrix) 的值越大. 因此如果我们不用 N 来标准化, 在梯度下降过程中第一层 (在池化层之前) 的损失计算就会过于重要. 我们当然不希望这样, 因为我们感兴趣的风格特征都在最深的那些层!

接着, 风格损失模块被以和内容损失模块相同的方式实现, 但是我们还得把 gramMatrix 加入作为参数:

  1. class StyleLoss(nn.Module):
  3.     def __init__(self, target, weight):
  4.         super(StyleLoss, self).__init__()
  5.         self.target = target.detach() * weight
  6.         self.weight = weight
  7.         self.gram = GramMatrix()
  8.         self.criterion = nn.MSELoss()
  10.     def forward(self, input):
  11.         self.output = input.clone()
  12.         self.G = self.gram(input)
  13.         self.G.mul_(self.weight)
  14.         self.loss = self.criterion(self.G, self.target)
  15.         return self.output
  17.     def backward(self, retain_graph=True):
  18.         self.loss.backward(retain_graph=retain_graph)
  19.         return self.loss

读取神经网络

现在, 我们要导入一个预训练好的神经网络. 和论文一样, 我们用预训练 的 19 层 VGG 网络 (VGG19).

PyTorch对 VGG 的实现模块分为两个子 Sequential 模块: features (包括卷积和池化层) 和 classifier (包括全连接层). 我们只对 features 感兴趣:

  1. cnn = models.vgg19(pretrained=True).features
  3. # 可能的话将它移到 GPU 上:
  4. if use_cuda:
  5.     cnn = cnn.cuda()

Sequential (顺序) 模块包含一个子模块的列表. 比如, vgg19.features 包含一个以正确深度排列的序列 (Conv2d, ReLU, Maxpool2d, Conv2d, ReLU…), 就如我们在 Content loss 部分讲到的, 我们想要把我们的风格和内容损失模块以想要的深度作为 ‘透明层’ 加入到 我们的网络中. 为了这样, 我们建立了一个新的 Sequential (顺序) 模块, 在其中我们把 vgg19 和我们的损失模块以正确的顺序加入:

  1. # 希望计算风格/内容损失的层 :
  2. content_layers_default = ['conv_4']
  3. style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
  5. def get_style_model_and_losses(cnn, style_img, content_img,
  6.                                style_weight=1000, content_weight=1,
  7.                                content_layers=content_layers_default,
  8.                                style_layers=style_layers_default):
  9.     cnn = copy.deepcopy(cnn)
  11.     # 仅为了有一个可迭代的列表 内容/风格 损失
  12.     content_losses = []
  13.     style_losses = []
  15.     model = nn.Sequential()  # 新建的 Sequential 网络模块
  16.     gram = GramMatrix()  # 我们需要一个克模块 (gram module) 来计算风格目标
  18.     # 可能的话将这些模块移到 GPU 上:
  19.     if use_cuda:
  20.         model = model.cuda()
  21.         gram = gram.cuda()
  23.     i = 1
  24.     for layer in list(cnn):
  25.         if isinstance(layer, nn.Conv2d):
  26.             name = "conv_" + str(i)
  27.             model.add_module(name, layer)
  29.             if name in content_layers:
  30.                 # 加内容损失:
  31.                 target = model(content_img).clone()
  32.                 content_loss = ContentLoss(target, content_weight)
  33.                 model.add_module("content_loss_" + str(i), content_loss)
  34.                 content_losses.append(content_loss)
  36.             if name in style_layers:
  37.                 # 加风格损失:
  38.                 target_feature = model(style_img).clone()
  39.                 target_feature_gram = gram(target_feature)
  40.                 style_loss = StyleLoss(target_feature_gram, style_weight)
  41.                 model.add_module("style_loss_" + str(i), style_loss)
  42.                 style_losses.append(style_loss)
  44.         if isinstance(layer, nn.ReLU):
  45.             name = "relu_" + str(i)
  46.             model.add_module(name, layer)
  48.             if name in content_layers:
  49.                 # 加内容损失:
  50.                 target = model(content_img).clone()
  51.                 content_loss = ContentLoss(target, content_weight)
  52.                 model.add_module("content_loss_" + str(i), content_loss)
  53.                 content_losses.append(content_loss)
  55.             if name in style_layers:
  56.                 # 加风格损失:
  57.                 target_feature = model(style_img).clone()
  58.                 target_feature_gram = gram(target_feature)
  59.                 style_loss = StyleLoss(target_feature_gram, style_weight)
  60.                 model.add_module("style_loss_" + str(i), style_loss)
  61.                 style_losses.append(style_loss)
  63.             i += 1
  65.         if isinstance(layer, nn.MaxPool2d):
  66.             name = "pool_" + str(i)
  67.             model.add_module(name, layer)  # ***
  69.     return model, style_losses, content_losses

注解：

在这篇论文中他们推荐将最大池化层更改为平均池化层. AlexNet是一个比 VGG19 更小的网络, 用它实现的话我们也不会看到 任何结果质量的不同. 而如果你想做这个替代的话, 可以用这些代码:

  1. # avgpool = nn.AvgPool2d(kernel_size=layer.kernel_size,
  2. #                         stride=layer.stride, padding = layer.padding)
  3. # model.add_module(name,avgpool)

输入图像

为了简化代码, 我们用与内容和风格图像同样尺寸的图像做输入. 这个图像可以是白噪声的, 也可以是一份内容图像的拷贝.

  1. input_img = content_img.clone()
  2. # 如果你想用白噪声做输入, 请取消下面的注释行:
  3. # input_img = Variable(torch.randn(content_img.data.size())).type(dtype)
  5. # 在绘图中加入原始的输入图像:
  6. plt.figure()
  7. imshow(input_img.data, title='Input Image')

梯度下降

由于本算法的作者 Leon Gatys 的建议 here, 我们将使用 L-BFGS 算法来跑我们的梯度下降. 和训练一个网络不同的是, 我们希望训练输入图像来最小化 内容/风格 损失. 我们想简单地建一个 PyTorch L-BFGS 优化器, 传入我们的图像作为变量进行优化. 但是 optim.LBFGS 的第一个形参是一个需要梯度的 PyTorch Variable . 我们的输入图像是一个 Variable , 但不是需要计算梯度的树的叶节点. 为了使这个变量需要梯度运算, 一个可能的方法是从输入图像构建一个 Parameter (参数) 对象. 然后我们只需给优化器的构造器传递一个 包含这个参数的列表:

  1. def get_input_param_optimizer(input_img):
  2.     # 这行显示了输入是一个需要梯度计算的参数
  3.     input_param = nn.Parameter(input_img.data)
  4.     optimizer = optim.LBFGS([input_param])
  5.     return input_param, optimizer

最后一步: 循环进行梯度下降. 每一步中我们必须喂给神经网络更新后 的输入以计算新的损失, 我们要运行每个损失的 backward 方法来动态 计算他们的梯度并呈现梯度下降的每一步. 这个优化器需要一个 “closure” : 一个重新评估模型并返回损失的函数.

然而, 这里有一个小问题. 被优化的图像的像素值会在 -\infty 和 +\infty 之间波动, 而不是继续保持在 0 到 1. 换句话说, 图像可能会被完美地优化成荒谬的值. 事实上, 我们必须在限制下使用 优化器来使我们的输入图像一直保持正确的值. 有一个简单的解决方案: 在每一步, 都校正图像使其保持 0-1 范围的值.

  1. def run_style_transfer(cnn, content_img, style_img, input_img, num_steps=300,
  2.                        style_weight=1000, content_weight=1):
  3.     """Run the style transfer."""
  4.     print('Building the style transfer model..')
  5.     model, style_losses, content_losses = get_style_model_and_losses(cnn,
  6.         style_img, content_img, style_weight, content_weight)
  7.     input_param, optimizer = get_input_param_optimizer(input_img)
  9.     print('Optimizing..')
  10.     run = [0]
  11.     while run[0] <= num_steps:
  13.         def closure():
  14.             # 校正更新后的输入图像值
  15.             input_param.data.clamp_(0, 1)
  17.             optimizer.zero_grad()
  18.             model(input_param)
  19.             style_score = 0
  20.             content_score = 0
  22.             for sl in style_losses:
  23.                 style_score += sl.backward()
  24.             for cl in content_losses:
  25.                 content_score += cl.backward()
  27.             run[0] += 1
  28.             if run[0] % 50 == 0:
  29.                 print("run {}:".format(run))
  30.                 print('Style Loss : {:4f} Content Loss: {:4f}'.format(
  31.                     style_score.data[0], content_score.data[0]))
  32.                 print()
  34.             return style_score + content_score
  36.         optimizer.step(closure)
  38.     # 最后一次的校正...
  39.     input_param.data.clamp_(0, 1)
  41.     return input_param.data

最后, 运行算法

  1. output = run_style_transfer(cnn, content_img, style_img, input_img)
  3. plt.figure()
  4. imshow(output, title='Output Image')
  6. # sphinx_gallery_thumbnail_number = 4
  7. plt.ioff()
  8. plt.show()