14.12. 神经风格迁移¶ 在 SageMaker Studio Lab 中打开 Notebook
如果你是一位摄影爱好者，你也许接触过滤镜。它能改变照片的颜色风格，从而使风景照更加锐利或者人像照皮肤更加美白。然而，一个滤镜通常只能改变照片的某个方面。如果要照片达到理想的风格，你可能需要尝试大量不同的滤镜组合。这个过程的复杂程度不亚于模型调参。
在本节中，我们将学习如何使用卷积神经网络，来自动将一个图像中的风格应用在另一图像之上，即*风格迁移*（style transfer） (Gatys et al., 2016)。 这项工作需要两张输入图像：一张是*内容图像*，另一张是*风格图像*。 我们将使用神经网络修改内容图像，使其在风格上更接近风格图像。 例如， 图 14.12.1中的内容图像是我们在西雅图郊区的瑞尼尔山国家公园拍摄的风景照，而风格图像则是一幅主题为秋天橡树的油画。 在输出的合成图像中，应用了风格图像的油画笔触，颜色也更加鲜艳，同时保留了内容图像中物体的主要形状。
图 14.12.1 给定内容图像和风格图像，风格迁移输出合成图像。¶
14.12.1. 方法¶
图 14.12.2用一个简化的例子说明了基于卷积神经网络的风格迁移方法。 首先，我们初始化合成图像，例如，将其初始化为内容图像。 该合成图像是风格迁移过程中唯一需要更新的变量，即风格迁移所需迭代的模型参数。 然后，我们选择一个预训练的卷积神经网络来抽取图像的特征，在模型训练过程中无须更新模型的参数。 这个深度卷积神经网络凭借多个层逐级抽取图像的特征，我们可以选择其中某些层的输出作为内容特征或风格特征。 以 图 14.12.2为例，这里选取的预训练的神经网络含有3个卷积层，其中第二层输出内容特征，第一层和第三层输出风格特征。
图 14.12.2 基于卷积神经网络的风格迁移。实线表示正向传播，虚线表示反向传播。¶
接下来，我们通过正向传播（实线箭头方向）计算风格迁移的损失函数，并通过反向传播（虚线箭头方向）迭代模型参数，即不断更新合成图像。 风格迁移常用的损失函数由3部分组成： （i）*内容损失*（content loss）使合成图像与内容图像在内容特征上接近； （ii）*风格损失*（style loss）使合成图像与风格图像在风格特征上接近； （iii）*全变分损失*（total variation loss）则有助于减少合成图像中的噪点。 最后，当模型训练结束时，我们输出风格迁移的模型参数，即得到最终的合成图像。
下面，我们将通过一个具体的实验来了解风格迁移的技术细节。
14.12.2. 读取内容和风格图像¶
首先，我们读取内容图像和风格图像。从打印出的坐标轴可以看出，它们的尺寸有所不同。
%matplotlib inline
import torch
import torchvision
from torch import nn
from d2l import torch as d2l
d2l.set_figsize()
content_img = d2l.Image.open('../img/rainier.jpg')
d2l.plt.imshow(content_img);
style_img = d2l.Image.open('../img/autumn-oak.jpg')
d2l.plt.imshow(style_img);
%matplotlib inline
from mxnet import autograd, gluon, image, init, np, npx
from mxnet.gluon import nn
from d2l import mxnet as d2l
npx.set_np()
d2l.set_figsize()
content_img = image.imread('../img/rainier.jpg')
d2l.plt.imshow(content_img.asnumpy());
[22:41:40] ../src/storage/storage.cc:196: Using Pooled (Naive) StorageManager for CPU
style_img = image.imread('../img/autumn-oak.jpg')
d2l.plt.imshow(style_img.asnumpy());
14.12.3. 预处理和后处理¶
下面，定义两个函数对图像进行预处理和后处理。 `preprocess`函数对输入图像在RGB三个通道分别做标准化，并将结果变换成卷积神经网络接受的输入格式。`postprocess`函数则将输出图像中的像素值还原回标准化之前的值。由于图像打印函数要求每个像素的浮点数值在0到1之间，我们对小于0和大于1的值分别取0和1。
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))
rgb_mean = np.array([0.485, 0.456, 0.406])
rgb_std = np.array([0.229, 0.224, 0.225])
def preprocess(img, image_shape):
img = image.imresize(img, *image_shape)
img = (img.astype('float32') / 255 - rgb_mean) / rgb_std
return np.expand_dims(img.transpose(2, 0, 1), axis=0)
def postprocess(img):
img = img[0].as_in_ctx(rgb_std.ctx)
return (img.transpose(1, 2, 0) * rgb_std + rgb_mean).clip(0, 1)
14.12.4. 抽取特征¶
我们使用基于ImageNet数据集预训练的VGG-19模型来抽取图像特征 (Gatys et al., 2016)。
pretrained_net = torchvision.models.vgg19(pretrained=True)
Downloading: "https://download.pytorch.org/models/vgg19-dcbb9e9d.pth" to /home/ci/.cache/torch/hub/checkpoints/vgg19-dcbb9e9d.pth
100%|██████████| 548M/548M [00:02<00:00, 213MB/s]
pretrained_net = gluon.model_zoo.vision.vgg19(pretrained=True)
Downloading /opt/mxnet/models/vgg19-ad2f660d.zipf5c89524-14c4-49c9-9bb9-460a4480a2b2 from https://apache-mxnet.s3-accelerate.dualstack.amazonaws.com/gluon/models/vgg19-ad2f660d.zip...
为了抽取图像的内容特征和风格特征，我们可以选择VGG网络中某些层的输出。 一般来说，越靠近输入层，越容易抽取图像的细节信息；反之，则越容易抽取图像的全局信息。 为了在合成图像中避免过度保留内容图像的细节，我们选择VGG较靠近输出的层，也称为*内容层*，来输出图像的内容特征。 我们还从VGG中选择不同层的输出来匹配局部和全局的风格，这些层也称为*风格层*。 正如在 8.2节中介绍的，VGG网络使用了5个卷积块。 实验中，我们选择第四卷积块的最后一个卷积层作为内容层，选择每个卷积块的第一个卷积层作为风格层。 通过打印`pretrained_net`实例，我们可以获取这些层的索引。
style_layers, content_layers = [0, 5, 10, 19, 28], [25]
style_layers, content_layers = [0, 5, 10, 19, 28], [25]
在使用VGG层抽取特征时，我们只需要用到VGG从输入层到最靠近输出层的内容层或风格层之间的所有层。下面构建一个新的网络`net`，它只保留需要用到的VGG的所有层。
net = nn.Sequential(*[pretrained_net.features[i] for i in
range(max(content_layers + style_layers) + 1)])
net = nn.Sequential()
for i in range(max(content_layers + style_layers) + 1):
net.add(pretrained_net.features[i])
给定输入`X`，如果我们简单地调用前向传播`net(X)`，只能获得最后一层的输出。由于我们还需要中间层的输出，因此这里需要逐层计算，并保留内容层和风格层的输出。
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
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
下面定义了两个函数：`get_contents`函数对内容图像抽取内容特征；`get_styles`函数对风格图像抽取风格特征。 因为在训练时无须改变预训练的VGG的模型参数，所以我们可以在训练开始之前就提取出内容图像的内容特征，以及风格图像的风格特征。 由于合成图像是风格迁移所需迭代的模型参数，我们只能在训练过程中通过调用`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
def get_contents(image_shape, device):
content_X = preprocess(content_img, image_shape).copyto(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).copyto(device)
_, styles_Y = extract_features(style_X, content_layers, style_layers)
return style_X, styles_Y
14.12.5. 定义损失函数¶
下面我们来描述风格迁移的损失函数。损失函数由内容损失、风格损失和全变分损失三部分组成。
14.12.5.1. 内容损失¶
与线性回归中的损失函数类似，内容损失通过平方误差函数衡量合成图像与内容图像在内容特征上的差异。平方误差函数的两个输入均为`extract_features`函数计算所得到的内容层的输出。
def content_loss(Y_hat, Y):
# We detach the target content from the tree used to dynamically compute
# the gradient: this is a stated value, not a variable. Otherwise the loss
# will throw an error.
return torch.square(Y_hat - Y.detach()).mean()
def content_loss(Y_hat, Y):
return np.square(Y_hat - Y).mean()
14.12.5.2. 风格损失¶
风格损失与内容损失类似，也通过平方误差函数衡量合成图像与风格图像在风格上的差异。为了表达风格层输出的风格，我们先通过`extract_features`函数计算风格层的输出。假设该输出的样本数为1，通道数为\(c\)，高和宽分别为\(h\)和\(w\)，我们可以将此输出转换为矩阵\(\mathbf{X}\)，其有\(c\)行和\(hw\)列。这个矩阵可以被看作是由\(c\)个长度为\(hw\)的向量\(\mathbf{x}_1, \ldots, \mathbf{x}_c\)组合而成的。其中，向量\(\mathbf{x}_i\)代表了通道\(i\)的风格特征。
在这些向量的*格拉姆矩阵*（Gram matrix）\(\mathbf{X}\mathbf{X}^\top \in \mathbb{R}^{c \times c}\)中，\(i\)行\(j\)列的元素\(x_{ij}\)即向量\(\mathbf{x}_i\)和\(\mathbf{x}_j\)的内积。它表达了通道\(i\)和通道\(j\)上风格特征的相关性。我们用这样的格拉姆矩阵来表达风格层输出的风格。请注意，当\(hw\)的值较大时，格拉姆矩阵中的元素容易出现较大的值。我们还注意到，格拉姆矩阵的高和宽都是通道数\(c\)。为了让风格损失不受这些值大小的影响，下面的`gram`函数将格拉姆矩阵除以了矩阵中元素的个数，即\(chw\)。
def gram(X):
num_channels, n = X.shape[1], X.numel() // X.shape[1]
X = X.reshape((num_channels, n))
return torch.matmul(X, X.T) / (num_channels * n)
def gram(X):
num_channels, n = X.shape[1], d2l.size(X) // X.shape[1]
X = X.reshape((num_channels, n))
return np.dot(X, X.T) / (num_channels * n)
很自然，风格损失的平方误差函数的两个格拉姆矩阵输入分别基于合成图像与风格图像的风格层输出。这里假设基于风格图像的格拉姆矩阵`gram_Y`已经预先计算好了。
def style_loss(Y_hat, gram_Y):
return torch.square(gram(Y_hat) - gram_Y.detach()).mean()
def style_loss(Y_hat, gram_Y):
return np.square(gram(Y_hat) - gram_Y).mean()
14.12.5.3. 全变分损失¶
有时候，我们学到的合成图像里面有大量高频噪点，即有特别亮或者特别暗的颗粒状像素。一种常用的降噪方法是*全变分降噪*（total variation denoising）。假设\(x_{i, j}\)表示坐标\((i, j)\)处的像素值，降低全变分损失
能够尽可能使邻近的像素值相似。
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())
def tv_loss(Y_hat):
return 0.5 * (np.abs(Y_hat[:, :, 1:, :] - Y_hat[:, :, :-1, :]).mean() +
np.abs(Y_hat[:, :, :, 1:] - Y_hat[:, :, :, :-1]).mean())
14.12.5.4. 损失函数¶
风格迁移的损失函数是内容损失、风格损失和全变分损失的加权和。通过调节这些权重超参数，我们可以权衡合成图像在保留内容、迁移风格以及降噪三方面的表现。
content_weight, style_weight, tv_weight = 1, 1e4, 10
def compute_loss(X, contents_Y_hat, styles_Y_hat, contents_Y, styles_Y_gram):
# Calculate the content, style, and total variance losses respectively
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
# Add up all the losses
l = sum(styles_l + contents_l + [tv_l])
return contents_l, styles_l, tv_l, l
content_weight, style_weight, tv_weight = 1, 1e4, 10
def compute_loss(X, contents_Y_hat, styles_Y_hat, contents_Y, styles_Y_gram):
# Calculate the content, style, and total variance losses respectively
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
# Add up all the losses
l = sum(styles_l + contents_l + [tv_l])
return contents_l, styles_l, tv_l, l
14.12.6. 初始化合成图像¶
在风格迁移中，合成的图像是训练期间唯一需要更新的变量。因此，我们可以定义一个简单的模型`SynthesizedImage`，并将合成的图像视为模型参数。模型的前向传播只需返回模型参数即可。
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
class SynthesizedImage(nn.Block):
def __init__(self, img_shape, **kwargs):
super(SynthesizedImage, self).__init__(**kwargs)
self.weight = self.params.get('weight', shape=img_shape)
def forward(self):
return self.weight.data()
下面，我们定义`get_inits`函数。该函数创建了合成图像的模型实例，并将其初始化为图像`X`。风格图像在各个风格层的格拉姆矩阵`styles_Y_gram`将在训练前预先计算好。
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
def get_inits(X, device, lr, styles_Y):
gen_img = SynthesizedImage(X.shape)
gen_img.initialize(init.Constant(X), ctx=device, force_reinit=True)
trainer = gluon.Trainer(gen_img.collect_params(), 'adam',
{'learning_rate': lr})
styles_Y_gram = [gram(Y) for Y in styles_Y]
return gen_img(), styles_Y_gram, trainer
14.12.7. 训练模型¶
在训练模型进行风格迁移时，我们不断抽取合成图像的内容特征和风格特征，然后计算损失函数。下面定义了训练循环。
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
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)
animator = d2l.Animator(xlabel='epoch', ylabel='loss',
xlim=[10, num_epochs], ylim=[0, 20],
legend=['content', 'style', 'TV'],
ncols=2, figsize=(7, 2.5))
for epoch in range(num_epochs):
with autograd.record():
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(1)
if (epoch + 1) % lr_decay_epoch == 0:
trainer.set_learning_rate(trainer.learning_rate * 0.8)
if (epoch + 1) % 10 == 0:
animator.axes[1].imshow(postprocess(X).asnumpy())
animator.add(epoch + 1, [float(sum(contents_l)),
float(sum(styles_l)), float(tv_l)])
return X
现在我们开始训练模型。我们将内容图像和风格图像的高和宽分别调整为300和450像素。我们用内容图像来初始化合成图像。
device, image_shape = d2l.try_gpu(), (300, 450) # PIL Image (h, w)
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)
device, image_shape = d2l.try_gpu(), (450, 300)
net.collect_params().reset_ctx(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.9, 500, 50)
我们可以看到，合成图像保留了内容图像的景物和物体，并同时迁移了风格图像的色彩。例如，合成图像具有与风格图像中一样的色块。其中一些色块甚至带有画笔的细微纹理。
14.12.8. 小结¶
风格迁移常用的损失函数由3部分组成：（i）内容损失使合成图像与内容图像在内容特征上接近；（ii）风格损失使合成图像与风格图像在风格特征上接近；（iii）全变分损失则有助于减少合成图像中的噪点。
我们可以通过预训练的卷积神经网络来抽取图像的特征，并通过最小化损失函数来不断更新合成图像，将其作为模型参数进行迭代。
我们用格拉姆矩阵来表示从风格层输出的风格。