实时风格迁移原来这么酷！用PyTorch分分钟搞定

雷克世界 · 公众号 · 机器人 · 2017-08-03 17:25

正文

原文来源：medium

作者：CeShine Lee

「机器人圈」编译：嗯~阿童木呀

在这篇文章中，我将结合自身经验，简要讲解如何用Pytorch编写并训练实时风格迁移模型。而这项研究主要是以Abhishek Kadian的实现为基础的，其运行效果非常完美。而我已经对它做了一些修改，一方面是为了让它变得更加有趣，另一方面是想让它能够更加适用于 Pytorch的运行环境。

该模型使用的是《基于感知损失的实时风格迁移和超分辨率》以及实例归一化（Instance Normalization）（未实施超分辨率）所描述那种方法。

下面就是我所添加到该模型实现中的三个主要部分：

1.使用官方预训练的VGG模型。

2.在训练期间输入中级训练结果。

3.添加论文中所描述的总变差正则化（Total Variation Regularization）。

使用官方预训练的VGG模型

论文《基于感知损失的实时风格迁移和超分辨率》中的模型架构

首先，我们需要快速浏览一下该模型架构。可以说该论文的主要贡献在于，它提出将生成的图像反馈到预先训练的图像分类模型，并从一些中间层中提取输出以计算损耗，便将会产生和Gatys等人所得到的类似的结果，但这明显具有更少的计算资源。因此，该结构的第一部分是 “图像变换网”（Image Transform Net），它可以从输入图像中生成新的图像。而第二部分只是一个“损失网络”，即前馈部分。而损失网络的权重是固定的，在训练过程中不会更新。

Abhishek的实现使用的是具有BGR信道顺序和中心信道偏移量为[-103.939，-116.779，-123.680]（它似乎也是论文中所描述使用的一种方法）的传统VGG模型。官方pytorch预训练模型使用的是一个统一格式：

所有预训练的模型都希望以相同的方式对输入图像进行归一化，即，预先形成小批量的3通道RGB图像（形式为3 x H x W），其中，H和W的预期值至少为224。图像必须加载到[0,1]的范围内，然后使用mean = [0.485, 0.456, 0.406]和std = [0.229, 0.224, 0.225]进行规范化。

以下是从官方预训练的模型中提取输出的代码：

启动：

除非明确指定，否则VGG模型中没有批量归一化（batch normalization）。所以，相较于之前的实现，该激活函数的值有很大的不同。一般来说，你需要放大风格损失（格拉姆矩阵）（gram matrix），因为大多数激活函数值小于1，而使用点积会使其更小。

在训练期间中级结果的输出

中级阶段，第75200个训练样本

当调整内容权重与风格权重比时，将会带来很大的帮助。你可以在训练期间停止训练，重新调整参数，而不必等待4个小时当完成训练之后才开始调整。

按照本文所述添加总变差正则化

该论文在实验部分提到了这一点——总变差正则化，但是似乎Abhishek没有实现这一目标：

输出图像是通过一种强度范围在1×10e^-6和1×10^e-4之间的总变差正则化进行正则化的，这是由对每一种风格目标进行交叉验证选择得来的。

维基百科：二维信号图像的总变差正则化公式

其实这是很容易实现的：

Pytorch autograd将为你处理反向传播（backward propagation）。在实际上，我还没有找到该如何以一种较为恰当的方式来调整正则化的权重。到目前为止，我使用的权重似乎在输出图像上并没有太大差异。

训练结果

该模型使用的是Microsoft COCO数据集进行训练的。图像的大小重新调整为256 x 256，网络在大约2个时期内进行训练，批量大小为4（与论文所述相同）。使用GTX1070进行训练的时间约为4至4.5小时，与论文报告中所述的用时大致相当。基于我那稍显粗略的实验来说，其中大量的时间用来对输入图像进行规范化处理。如果我们使用原始的VGG模型（未经测试），训练可能会进行得更快。在一些手动调整之后，内容权重vs风格比通常设置为1：10e^3〜10e^5。