精通 PyTorch——神经风格迁移(Neural Style Transfer)

501 阅读20分钟

在上一章中,我们开始使用 PyTorch 探索生成模型。通过在文本和音乐数据上进行无监督训练,我们构建了可以生成文本和音乐的机器学习模型。在本章中,我们将继续探索生成建模,并将类似的方法应用于图像数据。

我们将混合两幅不同图像 A 和 B 的不同方面,生成一幅包含图像 A 内容和图像 B 风格的结果图像 C。这个任务通常被称为神经风格迁移(Neural Style Transfer),因为在某种程度上,我们将图像 B 的风格转移到图像 A,以实现图像 C,如图 8.1 所示:

image.png

首先,我们将简要讨论如何解决这个问题,并理解实现风格迁移的基本思想。然后,我们将使用 PyTorch 实现我们自己的神经风格迁移系统,并将其应用于一对图像。在这个实现练习中,我们还将尝试理解风格迁移机制中不同参数的影响。

在本章结束时,你将理解神经风格迁移背后的概念,并能够使用 PyTorch 构建和测试你自己的神经风格迁移模型。

本章涵盖以下主题:

  • 理解如何在图像之间进行风格迁移
  • 使用 PyTorch 实现神经风格迁移

本章的所有代码文件可以在以下链接找到:github.com/arj7192/Mas…

现在让我们来看一下神经风格迁移的概念及其相关数学原理。

理解如何在图像之间进行风格迁移

在第2章《深度卷积神经网络架构》中,我们详细讨论了卷积神经网络(CNN)。CNN 在处理图像数据时,例如图像分类和目标检测任务中,是一些最成功的模型之一。其成功的核心原因之一是卷积层能够学习到空间表示。

例如,在一个识别狗和猫的分类器中,CNN 模型能够在提取高级特征的同时捕捉到图像的内容,这有助于它检测出与猫相比的狗的特定特征。我们将利用 CNN 图像分类器的这种能力来理解图像的内容。

我们知道 VGG 是一个强大的图像分类模型,如第2章《深度卷积神经网络架构》中讨论的那样。我们将使用 VGG 模型的卷积部分(不包括线性层)从图像中提取与内容相关的特征。

我们知道,每个卷积层会生成 N 个尺寸为 X*Y 的特征图。例如,假设我们有一个大小为 (3,3) 的单通道(灰度)输入图像,并且卷积层的输出通道数 (N) 是 3,卷积核大小是 (2,2),步幅为 (1,1),没有填充。这个卷积层将生成 3 个大小为 2x2 的特征图,因此在这种情况下,X=2,Y=2。

我们可以将这些由卷积层生成的 N 个特征图表示为一个大小为 NM 的二维矩阵,其中 M=XY。通过将每个卷积层的输出定义为二维矩阵,我们可以为每个卷积层定义一个损失函数。

这个损失函数称为内容损失,是卷积层期望输出与预测输出之间的平方损失,如图 8.2 所示,其中 N=3,X=2,Y=2:

image.png

正如我们所看到的,在这个示例中,输入图像(根据我们在图 8.1 中的标记,图像 C)通过卷积层(conv 层)被转换为三个特征图。这三个特征图的大小均为 2x2,然后格式化为一个 3x4 的矩阵。该矩阵与通过相同流程传递内容图像 A 所获得的期望输出进行比较,计算像素级的平方和损失,我们将其称为内容损失。

现在,为了从图像中提取风格,我们将使用从减少后的二维矩阵表示的行之间的内积生成的 Gram 矩阵【1】,如图 8.3 所示:

image.png

Gram矩阵携带了不同卷积特征图之间的相关性(在我们的图 8.3 示例中有三个特征图)。如果预测的Gram矩阵与期望的Gram矩阵接近,这意味着在这两种情况下(期望和预测),这三个卷积特征图之间具有相似的相关性(而非绝对值)。相关性表示的是像素之间的关系,而不是绝对的像素值。因此,Gram矩阵以图像的风格或纹理的形式捕捉了这种关系。

相比内容损失的计算,Gram矩阵计算是这里唯一额外的步骤。此外,如我们所见,像素级平方和损失的输出比内容损失要大得多。因此,这个数字通过除以NXY进行归一化,即特征图的数量(N)乘以特征图的长度(X)乘以宽度(Y)。这也有助于在不同卷积层之间标准化风格损失度量,因为这些层具有不同的N、X和Y。实现的详细信息可以在引入神经风格迁移的原始论文中找到【2】。

现在我们理解了内容损失和风格损失的概念,让我们来看看神经风格迁移是如何工作的,具体如下:

  1. 对于给定的VGG(或其他任何CNN)网络,我们定义网络中哪些卷积层应该附加内容损失。对风格损失进行相同的操作。
  2. 一旦我们有了这些列表,我们将内容图像通过网络,并计算在那些要计算内容损失的卷积层中的期望卷积输出(二维矩阵)。
  3. 接下来,我们将风格图像通过网络,并在那些要计算风格损失的卷积层中计算期望的Gram矩阵,如图8.4所示。

在下图中,内容损失将在第二和第三卷积层计算,而风格损失将在第二、第三和第五卷积层计算:

image.png

现在我们已经在选定的卷积层上得到了内容和风格的目标,我们可以生成一张包含内容图像内容和风格图像风格的图片了。

在初始化时,我们可以使用一个随机噪声矩阵作为生成图像的起点,或者直接使用内容图像作为起点。我们将这张图像传递到网络中,并在预选的卷积层上计算风格损失和内容损失。将所有风格损失相加,得到总风格损失;将所有内容损失相加,得到总内容损失。最后,通过加权的方式将这两部分损失相加,得到总损失。

如果我们给风格组件更多的权重,生成的图像将更能体现风格,反之亦然。使用梯度下降法,我们将损失反向传播回输入图像,从而更新生成图像。经过几个周期后,生成的图像应该会逐步演变,最终产生能够最小化各自损失的内容和风格表示,从而生成一个带有风格迁移效果的图像。

在图8.4中,池化层使用的是基于平均池化的方法,而不是传统的最大池化。为了确保梯度流动的平滑性,风格迁移中刻意使用了平均池化。我们希望生成的图像中像素之间没有过于明显的变化。此外,值得注意的是,上述图中的网络在最后一个计算风格或内容损失的层结束。因此,在这种情况下,因为在原始网络的第六卷积层没有计算损失,所以在风格迁移的上下文中讨论超过第五卷积层的内容是没有意义的。

在下一节中,我们将使用PyTorch实现我们自己的神经风格迁移系统。借助预训练的VGG模型,我们将使用本节讨论的概念生成具有艺术风格的图像。我们还将探索调整各种模型参数对生成图像内容和纹理/风格的影响。

使用 PyTorch 实现神经风格迁移

在我们讨论了神经风格迁移系统的内部工作原理后,现在我们可以使用 PyTorch 构建一个这样的系统了。作为一个练习,我们将加载一张风格图像和一张内容图像,然后加载预训练的 VGG 模型。在定义了在哪些层上计算风格和内容损失后,我们将修剪模型,只保留相关的层。最后,我们将训练神经风格迁移模型,通过不断迭代来优化生成的图像。

加载内容图像和风格图像

在这个练习中,我们只展示了代码中的关键部分以便演示。要访问完整的代码,请参见我们的 GitHub 仓库。请按照以下步骤操作:

首先,我们需要导入必要的库:

from PIL import Image
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision

dvc = torch.device("cuda" if torch.cuda.is_available() else "cpu")

除了其他库之外,我们还导入了 torchvision 库,用于加载预训练的 VGG 模型以及其他与计算机视觉相关的实用工具。

接下来,我们需要一张风格图像和一张内容图像。我们可以使用 Unsplash 网站 [4] 下载每种图像。下载的图像已经包含在本书的代码仓库中。在以下代码中,我们编写了一个函数,用于将图像加载为张量:

def image_to_tensor(image_filepath, image_dimension=128):
    img = Image.open(image_filepath).convert('RGB')
    # 显示图像
    ...
    torch_transformation = torchvision.transforms.Compose([
        torchvision.transforms.Resize(image_dimension),
        torchvision.transforms.ToTensor()])
    img = torch_transformation(img).unsqueeze(0)
    return img.to(dvc, torch.float)

然后,我们可以使用以下代码加载风格图像和内容图像:

style_image = image_to_tensor("./images/style.jpg")
content_image = image_to_tensor("./images/content.jpg")

unsqueeze 命令在张量的第 0 轴添加一个维度,将一个 (32, 32) 尺寸的张量转换为 (1, 32, 32) 尺寸的张量。你可以随意将参数从 0 改为 1,看看会发生什么。

上述代码应生成以下输出:

image.png

内容图像是泰姬陵的实景照片,而风格图像则是一幅艺术画作。通过风格迁移,我们希望生成一幅具有艺术风格的泰姬陵画作。然而,在此之前,我们需要加载并修剪 VGG19 模型。

加载和修剪预训练的 VGG19 模型

在本部分练习中,我们将使用预训练的 VGG 模型,并保留其卷积层。我们将对模型进行一些小的修改,使其可用于神经风格迁移。让我们开始吧:

我们首先加载预训练的 VGG19 模型,并使用其卷积层来生成内容和风格目标,从而分别得到内容和风格损失:

vgg19_model = torchvision.models.vgg19(pretrained=True).to(dvc)
print(vgg19_model)

输出应如下所示:

VGG(
(features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    ...
    (35): ReLU(inplace=True)
    (36): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    ...
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

我们不需要线性层,只需要模型的卷积部分。在上面的代码中,我们可以通过只保留模型对象的 features 属性来实现这一点:

vgg19_model = vgg19_model.features

在本练习中,我们不会调优 VGG 模型的参数。我们要调整的是生成图像的像素,而这些像素是在模型的输入端直接调整的。因此,我们需要确保加载的 VGG 模型的参数是固定的。

我们必须使用以下代码冻结 VGG 模型的参数:

for param in vgg19_model.parameters():
    param.requires_grad_(False)

现在,我们已经加载了 VGG 模型的相关部分,接下来我们需要将最大池化层更改为平均池化层,如前一节所述。在此过程中,我们将记录卷积层在模型中的位置:

conv_indices = []
for i in range(len(vgg19_model)):
    if vgg19_model[i]._get_name() == 'MaxPool2d':
        vgg19_model[i] = \
            nn.AvgPool2d(
                kernel_size=vgg19_model[i].kernel_size,
                stride=vgg19_model[i].stride,
                padding=vgg19_model[i].padding)
    if vgg19_model[i]._get_name() == 'Conv2d':
        conv_indices.append(i)
conv_indices = dict(enumerate(conv_indices, 1))
print(vgg19_model)

输出应如下所示:

Sequential(
  (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (1): ReLU(inplace=True)
  (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
...
  (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (35): ReLU(inplace=True)
  (36): AvgPool2d(kernel_size=2, stride=2, padding=0)
)

如我们所见,线性层已被移除,最大池化层已被替换为平均池化层,如突出显示的代码所示。

在上述步骤中,我们加载了一个预训练的 VGG 模型,并对其进行了修改,以便将其用作神经风格迁移模型。接下来,我们将把这个修改后的 VGG 模型转化为神经风格迁移模型。

构建神经风格迁移模型

此时,我们可以定义在哪些卷积层上计算内容和风格损失。在原始论文中,风格损失是在前五个卷积层上计算的,而内容损失则仅在第四个卷积层上计算。

我们将遵循相同的惯例,尽管你可以尝试不同的组合,并观察它们对生成图像的影响。请按照以下步骤操作:

首先,我们列出需要计算风格和内容损失的层:

layers = {1: 's', 2: 's', 3: 's', 4: 'sc', 5: 's'}

在这里,我们定义了从第一个到第五个卷积层,这些层与风格损失相关联,而第四个卷积层与内容损失相关联。注意,sc 分别表示风格和内容损失。

现在,让我们移除 VGG 模型中不必要的部分。我们只保留到第五个卷积层的内容,如下所示:

vgg_layers = nn.ModuleList(vgg19_model)
last_layer_idx = conv_indices[max(layers.keys())]
vgg_layers_trimmed = vgg_layers[:last_layer_idx+1]
neural_style_transfer_model = nn.Sequential(*vgg_layers_trimmed)
print(neural_style_transfer_model)

这应给我们以下输出:

Sequential( 
  (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (1): ReLU(inplace=True)
  (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  ...
  (8): ReLU(inplace=True)
  (9): AvgPool2d(kernel_size=2, stride=2, padding=0)  
  (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
)

正如我们所见,我们将具有 16 个卷积层的 VGG 模型转化为具有 5 个卷积层的神经风格迁移模型。

训练风格迁移模型

在本节中,我们将开始生成的图像的处理。我们可以通过多种方式初始化此图像,例如使用随机噪声图像或使用内容图像作为初始图像。目前,我们将从随机噪声开始。稍后我们还将看到使用内容图像作为起点如何影响结果。请按照以下步骤操作:

以下代码演示了如何使用随机数初始化一个 torch 张量的过程:

'''initialize as the content image
   ip_image = content_image.clone()
   initialize as random noise:'''
ip_image = torch.randn(content_image.data.size(), device=dvc)
plt.figure()
plt.imshow(ip_image.squeeze(0).cpu().detach().numpy().transpose(1,2,0).clip(0,1));

这应生成以下输出:

image.png

最后,我们可以开始模型训练循环了。首先,我们将定义训练的 epoch 数量,指定风格损失和内容损失的相对权重,并使用 Adam 优化器进行基于梯度下降的优化,学习率设为 0.1:

num_epochs=180
wt_style=1e6
wt_content=1
style_losses = []
content_losses = []
opt = optim.Adam([ip_image.requires_grad_()], lr=0.1)

在开始训练循环时,我们会在 epoch 的开始将风格和内容损失初始化为零,并将输入图像的像素值限制在 0 到 1 之间,以确保数值稳定性:

for curr_epoch in range(1, num_epochs+1):
    ip_image.data.clamp_(0, 1)
    opt.zero_grad()
    epoch_style_loss = 0
    epoch_content_loss = 0

此时,我们已经到达训练迭代中的一个关键步骤。在这里,我们必须计算每个预定义的风格和内容卷积层的风格和内容损失。

各个层的风格损失和内容损失将累加起来,得到当前 epoch 的总风格损失和总内容损失:

    for k in layers.keys():
        if 'c' in layers[k]:
            target = \
                neural_style_transfer_model[
                    :conv_indices[k]+1]
                    (content_image).detach()
            ip = \
                neural_style_transfer_model[
                    :conv_indices[k]+1](ip_image)
            epoch_content_loss += \
                torch.nn.functional.mse_loss(ip, target)
        if 's' in layers[k]:
            target = gram_matrix(
            neural_style_transfer_model[
                :conv_indices[k]+1]
                (style_image)).detach()
            ip = \
                gram_matrix(
                    neural_style_transfer_model[
                        :conv_indices[k]+1](ip_image))
            epoch_style_loss += \
                torch.nn.functional.mse_loss(ip, target)

如上代码所示,对于风格和内容损失,首先我们使用风格图像和内容图像计算风格和内容的目标值(即实际值)。我们使用 .detach() 方法将这些目标值标记为不可训练的固定目标值。接下来,我们基于生成的图像作为输入,计算各风格和内容层的预测输出。最后,我们计算风格和内容损失。

对于风格损失,我们还需要使用预定义的 gram 矩阵函数计算 gram 矩阵,如下代码所示:

def gram_matrix(ip):
    num_batch, num_channels, height, width = ip.size()
    feats = ip.view(num_batch * num_channels, width * height)
    gram_mat = torch.mm(feats, feats.t())
    return gram_mat.div(
        num_batch * num_channels * width * height)

如前所述,我们可以使用 torch.mm 函数计算内积。这将计算 gram 矩阵,并通过将其除以特征图的数量、特征图的宽度和高度来对矩阵进行归一化。

继续在我们的训练循环中,现在我们已经计算出总的风格和内容损失,我们需要使用之前定义的权重,将这两者的加权和作为最终的总损失来计算:

    epoch_style_loss *= wt_style
    epoch_content_loss *= wt_content
    total_loss = epoch_style_loss + epoch_content_loss
    total_loss.backward()

最后,每隔 k 个 epoch,我们可以通过查看损失值以及生成的图像来观察训练的进展。以下图显示了前述代码在总共 180 个 epoch 中,每隔 20 个 epoch 记录一次的风格迁移生成图像的演变过程:

image.png

很明显,模型最初将风格图像的风格应用到随机噪声上。随着训练的进行,内容损失开始发挥作用,从而将内容赋予风格化的图像。到第 180 个 epoch 时,我们可以看到生成的图像,已经看起来像是一幅很好的泰姬陵艺术绘画。下图显示了从第 0 到第 180 个 epoch,风格损失和内容损失的逐渐减少情况:

image.png

显著的是,风格损失在初期急剧下降,这也在图 8.7 中得以体现,初始 epochs 中风格对图像的影响远大于内容。在训练的后期阶段,两者的损失都逐渐一起下降,最终得到的风格迁移图像在风格图像的艺术性和相片的现实主义之间达成了一个不错的折衷。

实验风格迁移系统

在前一节成功训练了风格迁移系统之后,我们现在来看一下该系统在不同超参数设置下的反应。按照以下步骤进行:

在前一节中,我们将内容权重设置为1,风格权重设置为1e6。现在我们将风格权重进一步增加10倍,即到1e7,并观察它对风格迁移过程的影响。使用新权重训练600个epochs后,我们得到以下风格迁移的进展情况:

image.png

在这里,我们可以看到,与之前的情况相比,要达到合理的结果需要更多的训练 epochs。更重要的是,较高的风格权重确实对生成的图像产生了影响。当我们将前面的图像与图 8.7 中的图像进行比较时,发现前者与图 8.5 中的风格图像有更强的相似性。

同样地,将风格权重从 1e6 降低到 1e5 会产生更注重内容的结果,如图 8.10 所示:

image.png

与较高风格权重的情况相比,较低的风格权重意味着在更少的训练 epochs 内就能获得看起来合理的结果。生成图像中的风格成分较少,主要由内容图像数据填充。在这种情况下,我们只训练了 6 个 epochs,因为在那之后结果基本趋于稳定。

最后一种变化可以是将生成图像初始化为内容图像,而不是随机噪声,同时保持原始的风格和内容权重为 1e6 和 1。图 8.11 显示了在这种情况下每个 epoch 的进展情况:

image.png

通过将上图与图 8.7 进行比较,我们可以看到,将内容图像作为起点为我们提供了不同的进展路径,以生成合理的风格迁移图像。似乎内容和风格成分更同时地施加在生成图像上,而不像图 8.7 中那样,先施加风格,然后再添加内容。以下图表验证了这一假设:

image.png

正如我们所见,随着训练的进行,风格损失和内容损失一起下降,并在最后趋于平稳。然而,无论是在图 8.17 和 8.11 以及图 8.9 和 8.10 中,最终的结果都表现为合理的泰姬陵艺术印象。图 8.12 中的两条损失线相比于图 8.8 中的显得更为分离,这只是因为两条曲线的 y 轴范围不同。

我们已经成功地使用 PyTorch 构建了一个神经风格迁移模型,通过使用一张内容图像(美丽的泰姬陵照片)和一张风格图像(画布画作),生成了一个合理的泰姬陵艺术画作的近似效果。这个应用可以扩展到其他各种组合上。交换内容图像和风格图像可能也会产生有趣的结果,并提供更多关于模型内部工作原理的洞察。

我们鼓励你通过以下方式扩展我们在本章中讨论的练习:

  • 更改风格和内容层的列表
  • 使用更大的图像尺寸
  • 尝试更多风格和内容损失权重的组合
  • 使用其他优化器,例如 SGD 和 LBFGS
  • 在不同的学习率下训练更长的 epoch,以观察这些方法产生的图像之间的差异

总结

在本章中,我们通过生成一张包含一张图像的内容和另一张图像风格的图片,将生成式机器学习的概念应用到了图像上,这个任务被称为神经风格迁移。

在下一章中,我们将扩展这一范式,研究一个生成器生成假数据并且一个判别器区分假数据和真实数据的模型。这类模型被广泛称为生成对抗网络 (GANs)。我们将在下一章中探索深度卷积 GANs (DCGANs)。