Pytorch 风格迁移（Style transfer）

Pytorch 风格迁移

0. 环境介绍

环境使用 Kaggle 里免费建立的 Notebook

教程使用李沐老师的 动手学深度学习 网站和 视频讲解

小技巧：当遇到函数看不懂的时候可以按 Shift+Tab 查看函数详解。

1. 风格迁移

1.1 概述

将一个图像中的风格应用在另一图像之上，即风格迁移（style transfer）。这里我们需要两张输入图像：一张是内容图像，另一张是风格图像。 我们将使用神经网络修改内容图像，使其在风格上接近风格图像。

1.2 基于 CNN 的样式迁移

首先初始化合成图像，例如可以使用内容图像初始化。
合成图像是训练过程中唯一需要更新的变量（不像 CNN 训练的是卷积层的权重）。

然后选择一个预训练的 CNN 来抽取图像特征（如 VGG等），其中模型参数在训练中不需要更新。
这个选择的 CNN 凭借多个层逐级抽取图像的特征，我们可以选择其中某些层的输出作为内容特征或风格特征。

这里选取的预训练的神经网络含有 $3$ 个卷积层，其中第二层输出内容特征，第一层和第三层输出风格特征。

接下来，我们通过前向传播（实线箭头方向）计算风格迁移的损失函数，并通过反向传播（虚线箭头方向）迭代模型参数，即不断更新合成图像。

风格迁移常用的损失函数由3部分组成：
（i）内容损失使合成图像与内容图像在内容特征上接近；
（ii）风格损失使合成图像与风格图像在风格特征上接近；
（iii）全变分损失（total variation denoising）则有助于减少合成图像中的噪点。

最后，当模型训练结束时，我们输出风格迁移的模型参数，即得到最终的合成图像。

2. 代码

2.1 导入图片

从 github 下载图片到本地：

!pip install -U d2l
%matplotlib inline
import torch
import os
import requests
import torchvision
from torch import nn
from d2l import torch as d2l

if not os.path.exists('../data'):
    os.mkdir('../data')

for name in ['rainier', 'autumn-oak']:
    r = requests.get('https://raw.githubusercontent.com/d2l-ai/d2l-zh/master/img/%s.jpg' % name) 
    with open('../data/%s.jpg' % name, 'wb') as f:
        f.write(r.content)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

读取内容图像：

d2l.set_figsize()
content_img = d2l.Image.open('../data/rainier.jpg')
d2l.plt.imshow(content_img)
1
2
3

读取风格图像：

style_img = d2l.Image.open('../data/autumn-oak.jpg')
d2l.plt.imshow(style_img)
1
2

2.2 预处理和后处理

预处理函数 preprocess 对输入图像在 RGB 三个通道分别做标准化，并将结果变换成卷积神经网络接受的输入格式（图片==>Tensor）。
后处理函数 postprocess 则将输出图像中的像素值还原回标准化之前的值。 由于图像打印函数要求每个像素的浮点数值在 $0$ 到 $1$ 之间，我们对小于 $0$ 和大于 $1$ 的值分别取 $0$ 和 $1$（Tensor==>图片）。

# imagenet 输入要求
rgb_mean = torch.tensor([0.485, 0.456, 0.406])
rgb_std = torch.tensor([0.229, 0.224, 0.225])

def preprocess(img, image_shape):
    transforms = torchvision.transforms.Compose([
        torchvision.transforms.Resize(image_shape),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(mean=rgb_mean, std=rgb_std)])
    return transforms(img).unsqueeze(0)

def postprocess(img):
    img = img[0].to(rgb_std.device)
    img = torch.clamp(img.permute(1, 2, 0) * rgb_std + rgb_mean, 0, 1)
    return torchvision.transforms.ToPILImage()(img.permute(2, 0, 1))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

2.3 抽取图像特征

使用基于ImageNet 数据集预训练的 VGG-19 模型来抽取图像特征：

pretrained_net = torchvision.models.vgg19(pretrained=True)
1

使用不同的层抽取风格和内容，其中越小越靠近输入（匹配局部信息），越大越靠近输出（匹配全局信息）。风格层均匀选取（既要局部，也要全局），内容层取靠后的（只考虑全局）：

style_layers, content_layers = [0, 5, 10, 19, 28], [25]
1

拿出网络，取到 $28$ 层，后面的就不要了：

net = nn.Sequential(*[pretrained_net.features[i] for i in
                      range(max(content_layers + style_layers) + 1)])
1
2

抽取特征：

def extract_features(X, content_layers, style_layers):
    contents = []
    styles = []
    for i in range(len(net)):
        X = net[i](X)
        if i in style_layers:
            styles.append(X)
        if i in content_layers:
            contents.append(X)
    return contents, styles
1
2
3
4
5
6
7
8
9
10

拿出内容特征 Tensor，拿出风格特征 Tensor 都调用了 extract_features：

def get_contents(image_shape, device):
    content_X = preprocess(content_img, image_shape).to(device)
    contents_Y, _ = extract_features(content_X, content_layers, style_layers)
    return content_X, contents_Y

def get_styles(image_shape, device):
    style_X = preprocess(style_img, image_shape).to(device)
    _, styles_Y = extract_features(style_X, content_layers, style_layers)
    return style_X, styles_Y
1
2
3
4
5
6
7
8
9

2.4 定义损失函数

2.4.1 内容损失

与线性回归中的损失函数类似，内容损失通过平方误差函数衡量合成图像与内容图像在内容特征上的差异。 平方误差函数的两个输入均为 extract_features 函数计算所得到的内容层的输出：

def content_loss(Y_hat, Y):
    # 我们从动态计算梯度的树中分离目标：
    # 这是一个规定的值，而不是一个变量。
    return torch.square(Y_hat - Y.detach()).mean()
1
2
3
4

2.4.2 风格损失

风格比较不好计算损失，风格怎么理解是否一致？两张图像风格一致并不要求像素值一样，只需要图像的统计分布相似即可。

def gram(X):
	# n 是每个通道像素数量 = h * w
    num_channels, n = X.shape[1], X.numel() // X.shape[1]
    X = X.reshape((num_channels, n))
    # X * X.T (3*hw × hw*3 = 3*3) 
    return torch.matmul(X, X.T) / (num_channels * n)
1
2
3
4
5
6

def style_loss(Y_hat, gram_Y):
    return torch.square(gram(Y_hat) - gram_Y.detach()).mean()
1
2

2.4.3 全变分损失

我们学到的合成图像里面有大量高频噪点，即有特别亮或者特别暗的颗粒像素。 一种常见的去噪方法是全变分去噪（total variation denoising）： 假设 $x_{i,j}$ 表示坐标 $(i,j)$处的像素值，降低全变分损失
$$\sum _{i,j}\left|x_{i,j}-x_{i+1,j}\right|+\left|x_{i,j}-x_{i,j+1}\right|$$
能够尽可能使邻近的像素值相似：

def tv_loss(Y_hat):
    return 0.5 * (torch.abs(Y_hat[:, :, 1:, :] - Y_hat[:, :, :-1, :]).mean() +
                  torch.abs(Y_hat[:, :, :, 1:] - Y_hat[:, :, :, :-1]).mean())
1
2
3

2.4.4 计算损失函数

风格转移的损失函数是内容损失、风格损失和总变化损失的加权和。
通过调节这些权重超参数，我们可以权衡合成图像在保留内容、迁移风格以及去噪三方面的相对重要性：

content_weight, style_weight, tv_weight = 1, 1e3, 10

def compute_loss(X, contents_Y_hat, styles_Y_hat, contents_Y, styles_Y_gram):
    # 分别计算内容损失、风格损失和全变分损失
    contents_l = [content_loss(Y_hat, Y) * content_weight for Y_hat, Y in zip(
        contents_Y_hat, contents_Y)]
    styles_l = [style_loss(Y_hat, Y) * style_weight for Y_hat, Y in zip(
        styles_Y_hat, styles_Y_gram)]
    tv_l = tv_loss(X) * tv_weight
    # 对所有损失求和
    l = sum(10 * styles_l + contents_l + [tv_l])
    return contents_l, styles_l, tv_l, l
1
2
3
4
5
6
7
8
9
10
11
12

2.5 初始化合成图像

在风格迁移中，合成的图像是训练期间唯一需要更新的变量，将合成的图像视为模型参数，模型的前向传播只需返回模型参数即可（可以直接使用内容图像作为合成图像的初始化）：

class SynthesizedImage(nn.Module):
    def __init__(self, img_shape, **kwargs):
        super(SynthesizedImage, self).__init__(**kwargs)
        self.weight = nn.Parameter(torch.rand(*img_shape))

    def forward(self):
        return self.weight
1
2
3
4
5
6
7

def get_inits(X, device, lr, styles_Y):
    gen_img = SynthesizedImage(X.shape).to(device)
    gen_img.weight.data.copy_(X.data)
    trainer = torch.optim.Adam(gen_img.parameters(), lr=lr)
    styles_Y_gram = [gram(Y) for Y in styles_Y]
    return gen_img(), styles_Y_gram, trainer
1
2
3
4
5
6

2.6 训练

训练模型进行风格迁移时，我们不断抽取合成图像的内容特征和风格特征，然后计算损失函数：

def train(X, contents_Y, styles_Y, device, lr, num_epochs, lr_decay_epoch):
    X, styles_Y_gram, trainer = get_inits(X, device, lr, styles_Y)
    scheduler = torch.optim.lr_scheduler.StepLR(trainer, lr_decay_epoch, 0.8)
    animator = d2l.Animator(xlabel='epoch', ylabel='loss',
                            xlim=[10, num_epochs],
                            legend=['content', 'style', 'TV'],
                            ncols=2, figsize=(7, 2.5))
    for epoch in range(num_epochs):
        trainer.zero_grad()
        contents_Y_hat, styles_Y_hat = extract_features(
            X, content_layers, style_layers)
        contents_l, styles_l, tv_l, l = compute_loss(
            X, contents_Y_hat, styles_Y_hat, contents_Y, styles_Y_gram)
        l.backward()
        trainer.step()
        scheduler.step()
        if (epoch + 1) % 10 == 0:
            animator.axes[1].imshow(postprocess(X))
            animator.add(epoch + 1, [float(sum(contents_l)),
                                     float(sum(styles_l)), float(tv_l)])
    return X
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

首先将内容图像和风格图像的高和宽分别调整为 $300$ 和 $450$ 像素，用内容图像来初始化合成图像：

device, image_shape = d2l.try_gpu(), (300, 450)
net = net.to(device)
content_X, contents_Y = get_contents(image_shape, device)
_, styles_Y = get_styles(image_shape, device)
output = train(content_X, contents_Y, styles_Y, device, 0.3, 500, 50)
1
2
3
4
5

设置稍大一点的像素看看效果：

device, image_shape = d2l.try_gpu(), (900, 1200)
net = net.to(device)
content_X, contents_Y = get_contents(image_shape, device)
_, styles_Y = get_styles(image_shape, device)
output = train(content_X, contents_Y, styles_Y, device, 0.3, 500, 50)
1
2
3
4
5