StyleTransferTrilogy

风格迁移三部曲

风格迁移是一个很有意思的任务，通过风格迁移可以使一张图片保持本身内容大致不变的情况下呈现出另外一张图片的风格。本文会介绍以下三种风格迁移方式以及对应的代码实现：

固定风格固定内容的普通风格迁移（A Neural Algorithm of Artistic Style）
固定风格任意内容的快速风格迁移（Perceptual Losses for Real-Time Style Transfer and Super-Resolution）
任意风格任意内容的极速风格迁移（Meta Networks for Neural Style Transfer）

本文所使用的环境是 pytorch 0.4.0，如果你使用了其他的版本，稍作修改即可正确运行。

固定风格固定内容的普通风格迁移

最早的风格迁移就是在固定风格、固定内容的情况下做的风格迁移，这是最慢的方法，也是最经典的方法。

最原始的风格迁移的思路很简单，把图片当做可以训练的变量，通过优化图片来降低与内容图片的内容差异以及降低与风格图片的风格差异，迭代训练多次以后，生成的图片就会与内容图片的内容一致，同时也会与风格图片的风格一致。

VGG16

VGG16 是一个很经典的模型，它通过堆叠 3x3 的卷积层和池化层，在 ImageNet 上获得了不错的成绩。我们使用在 ImageNet 上经过预训练的 VGG16 模型可以对图像提取出有用的特征，这些特征可以帮助我们去衡量两个图像的内容差异和风格差异。

在进行风格迁移任务时，我们只需要提取其中几个比较重要的层，所以我们对 pytorch 自带的预训练 VGG16 模型稍作了一些修改：

class VGG(nn.Module):

    def __init__(self, features):
        super(VGG, self).__init__()
        self.features = features
        self.layer_name_mapping = {
            '3': "relu1_2",
            '8': "relu2_2",
            '15': "relu3_3",
            '22': "relu4_3"
        }
        for p in self.parameters():
            p.requires_grad = False
    
    def forward(self, x):
        outs = []
        for name, module in self.features._modules.items():
            x = module(x)
            if name in self.layer_name_mapping:
                outs.append(x)
        return outs

vgg16 = models.vgg16(pretrained=True)
vgg16 = VGG(vgg16.features[:23]).to(device).eval()

经过修改的 VGG16 可以输出 relu1_2,relu2_2,relu3_3,relu4_3 这几个特定层的特征图。下面这两句代码就是它的用法：

features = vgg16(input_img)
content_features = vgg16(content_img)

举个例子，当我们使用 vgg16 对 input_img 计算特征时，它会返回四个矩阵给 features，假设 input_img 的尺寸是 [1, 3, 512, 512]（四个维度分别代表 batch, channels, height, width），那么它返回的四个矩阵的尺寸就是这样的：

relu1_2 [1, 64, 512, 512]
relu2_2 [1, 128, 256, 256]
relu3_3 [1, 256, 128, 128]
relu4_3 [1, 512, 64, 64]

内容

我们进行风格迁移的时候，必须保证生成的图像与内容图像的内容一致性，不然风格迁移就变成艺术创作了。那么如何衡量两张图片的内容差异呢？很简单，通过 VGG16 输出的特征图来衡量图片的内容差异。

提示：在本方法中没有 Image Transform Net，为了表述方便，我们使用了第二篇论文中的图。

这里使用的损失函数是：

$\Large\ell^{\phi,j}_{feat}(\hat{y},y)=\frac{1}{C_jH_jW_j}||\phi_j(\hat{y})-\phi_j(y)||^2_2$

其中：

$\hat{y}$ 是输入图像（也就是生成的图像）
$y$ 是内容图像
$\phi$ 代表 VGG16
$\j$ 在这里是 relu3_3
$\phi_j(x)$ 指的是 x 图像输入到 VGG 以后的第 j 层的特征图
$C_j\times H_j\times W_j$ 是第 j 层输出的特征图的尺寸

根据生成图像和内容图像在 $\text{relu3_3}$ 输出的特征图的均方误差（MeanSquaredError）来优化生成的图像与内容图像之间的内容一致性。

那么写成代码就是这样的：

content_loss = F.mse_loss(features[2], content_features[2]) * content_weight

因为我们这里使用的是经过在 ImageNet 预训练过的 VGG16 提取的特征图，所以它能提取出图像的高级特征，通过优化生成图像和内容图像特征图的 mse，可以迫使生成图像的内容与内容图像在 VGG16 的 relu3_3 上输出相似的结果，因此生成图像和内容图像在内容上是一致的。

风格

Gram 矩阵

那么如何衡量输入图像与风格图像之间的内容差异呢？这里就需要提出一个新的公式，Gram 矩阵：

$\Large{G^\phi_j(x)_{c,c'}=\frac{1}{C_jH_jW_j} \sum_{h=1}^{H_j} \sum_{w=1}^{W_j} \phi_j(x)_{h,w,c}\phi_j(x)_{h,w,c'}}$

其中：

$\hat{y}$ 是输入图像（也就是生成的图像）
$y$ 是风格图像
$C_j\times H_j\times W_j$ 是第 j 层输出的特征图的尺寸。
$G^\phi_j(x)$ 指的是 x 图像的第 j 层特征图对应的 Gram 矩阵，比如 64 个卷积核对应的卷积层输出的特征图的 Gram 矩阵的尺寸是 $(64,64)$ 。
$G^\phi_j(x)_{c,c'}$ 指的是 Gram 矩阵第 $(c,c')$ 坐标对应的值。
$\phi_j(x)$ 指的是 x 图像输入到 VGG 以后的第 j 层的特征图， $\phi_j(x)_{h,w,c}$ 指的是特征图 $(h,w,c)$ 坐标对应的值。

Gram 矩阵的计算方法其实很简单，Gram 矩阵的 $(c,c')$ 坐标对应的值，就是特征图的第 $c$ 张和第 $c'$ 张图对应元素相乘，然后全部加起来并且除以 $C_j\times H_j\times W_j$ 的结果。根据公式我们可以很容易推断出 Gram 矩阵是对称矩阵。

具体到代码，我们可以写出下面的函数：

def gram_matrix(y):
    (b, ch, h, w) = y.size()
    features = y.view(b, ch, w * h)
    features_t = features.transpose(1, 2)
    gram = features.bmm(features_t) / (ch * h * w)
    return gram

参考链接：

https://github.com/pytorch/examples/blob/0.4/fast_neural_style/neural_style/utils.py#L21-L26

假设我们输入了一个 [1, 3, 512, 512] 的图像，下面就是各个矩阵的尺寸：

relu1_2 [1, 64, 512, 512]，gram [1, 64, 64]
relu2_2 [1, 128, 256, 256]，gram [1, 128, 128]
relu3_3 [1, 256, 128, 128]，gram [1, 256, 256]
relu4_3 [1, 512, 64, 64]，gram [1, 512, 512]

风格损失

根据生成图像和风格图像在 relu1_2、relu2_2、relu3_3、relu4_3 输出的特征图的 Gram 矩阵之间的均方误差（MeanSquaredError）来优化生成的图像与风格图像之间的风格差异：

$\Large\ell^{\phi,j}_{style}(\hat{y},y)=||G^\phi_j(\hat{y})-G^\phi_j(y)||^2_F$

其中：

$\hat{y}$ 是输入图像（也就是生成的图像）
$y$ 是风格图像
$G^\phi_j(x)$ 指的是 x 图像的第 j 层特征图对应的 Gram 矩阵

那么写成代码就是下面这样：

style_grams = [gram_matrix(x) for x in style_features]

style_loss = 0
grams = [gram_matrix(x) for x in features]
for a, b in zip(grams, style_grams):
    style_loss += F.mse_loss(a, b) * style_weight

训练

那么风格迁移的目标就很简单了，直接将两个 loss 按权值加起来，然后对图片优化 loss，即可优化出既有内容图像的内容，也有风格图像的风格的图片。代码如下：

input_img = content_img.clone()
optimizer = optim.LBFGS([input_img.requires_grad_()])
style_weight = 1e6
content_weight = 1

run = [0]
while run[0] <= 300:
    def f():
        optimizer.zero_grad()
        features = vgg16(input_img)
        
        content_loss = F.mse_loss(features[2], content_features[2]) * content_weight
        style_loss = 0
        grams = [gram_matrix(x) for x in features]
        for a, b in zip(grams, style_grams):
            style_loss += F.mse_loss(a, b) * style_weight
        
        loss = style_loss + content_loss
        
        if run[0] % 50 == 0:
            print('Step {}: Style Loss: {:4f} Content Loss: {:4f}'.format(
                run[0], style_loss.item(), content_loss.item()))
        run[0] += 1
        
        loss.backward()
        return loss
    
    optimizer.step(f)

此处使用了 LBFGS，所以 loss 需要包装在一个函数里，代码参考了： https://pytorch.org/tutorials/advanced/neural_style_tutorial.html

效果

最终效果如图所示：

可以看到生成的图像既有风格图像的风格，也有内容图像的内容，很完美。不过生成一幅256x256 的图像在 1080ti 上需要18.6s，这个时间挺长的，谈不上实时性，因此我们可以来看看第二篇论文中的方法。

固定风格任意内容的快速风格迁移

有了上面的铺垫，理解固定风格任意内容的快速风格迁移就简单很多了。思路很简单，就是先搭建一个转换网络，然后通过优化转换网络的权值来实现快速风格迁移。由于这个转换网络可以接受任意图像，所以这是任意内容的风格迁移。

模型

模型结构很简单，分为三个部分：

降维，三层卷积层，逐渐提升通道数为128，并且通过 stride 把特征图的宽高缩小为原来的八分之一
5个 ResidualBlock 堆叠
升维，三层卷积层，逐渐降低通道数为3，并且通过 nn.Upsample 把特征图的宽高还原为原来的大小

先降维再升维是为了减少计算量，中间的 5 个 Residual 结构可以学习如何在原图上添加少量内容，改变原图的风格。下面让我们来看看代码。

ConvLayer

def ConvLayer(in_channels, out_channels, kernel_size=3, stride=1, 
    upsample=None, instance_norm=True, relu=True):
    layers = []
    if upsample:
        layers.append(nn.Upsample(mode='nearest', scale_factor=upsample))
    layers.append(nn.ReflectionPad2d(kernel_size // 2))
    layers.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride))
    if instance_norm:
        layers.append(nn.InstanceNorm2d(out_channels))
    if relu:
        layers.append(nn.ReLU())
    return layers

首先我们实现了一个函数，ConvLayer，它包含：

nn.Upsample（可选）
nn.ReflectionPad2d
nn.Conv2d
nn.InstanceNorm2d（可选）
nn.ReLU（可选）

因为每个卷积层前后都可能会用到这些层，为了简化代码，我们将它写成一个函数，返回这些层用于搭建模型。

ResidualBlock

class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super(ResidualBlock, self).__init__()
        self.conv = nn.Sequential(
            *ConvLayer(channels, channels, kernel_size=3, stride=1), 
            *ConvLayer(channels, channels, kernel_size=3, stride=1, relu=False)
        )

    def forward(self, x):
        return self.conv(x) + x

这里写的就不是函数，而是一个类，因为它内部包含许多层，而且并不是简单的自上而下的结构（Sequential），而是有了跨层的连接（self.conv(x) + x），所以我们需要继承 nn.Module，实现 forward 函数，才能实现跨层连接。

TransformNet

最后这个模型就很简单了，照着论文里给出的表格搭建即可。我们这里为了实验方便，添加了 base 参数，当 base=8 时，卷积核的个数是按 8, 16, 32 递增的，当 base=32 时，卷积核个数是按 32, 64, 128 递增的。有了这个参数，我们可以按需增加模型规模，base 越大，图像质量越好。

class TransformNet(nn.Module):
    def __init__(self, base=32):
        super(TransformNet, self).__init__()
        self.downsampling = nn.Sequential(
            *ConvLayer(3, base, kernel_size=9), 
            *ConvLayer(base, base*2, kernel_size=3, stride=2), 
            *ConvLayer(base*2, base*4, kernel_size=3, stride=2), 
        )
        self.residuals = nn.Sequential(*[ResidualBlock(base*4) for i in range(5)])
        self.upsampling = nn.Sequential(
            *ConvLayer(base*4, base*2, kernel_size=3, upsample=2),
            *ConvLayer(base*2, base, kernel_size=3, upsample=2),
            *ConvLayer(base, 3, kernel_size=9, instance_norm=False, relu=False),
        )
    
    def

StyleTransferTrilogy

Install / Use

README

StyleTransferTrilogy

风格迁移三部曲

固定风格固定内容的普通风格迁移

VGG16

内容

风格

Gram 矩阵

风格损失

训练

效果

固定风格任意内容的快速风格迁移

模型

ConvLayer

ResidualBlock

TransformNet