注意
点击此处下载完整的示例代码
使用 PyTorch 进行神经风格迁移¶
创建于: 2017年4月6日 | 最后更新于: 2024年1月19日 | 最后验证于: 2024年11月5日
作者: Alexis Jacq
编辑: Winston Herring
引言¶
本教程解释了如何实现 Leon A. Gatys、Alexander S. Ecker 和 Matthias Bethge 开发的 神经风格算法。神经风格，或称神经迁移，允许你获取一张图像并以新的艺术风格进行重现。该算法需要三张图像：一张输入图像、一张内容图像和一张风格图像，然后改变输入图像，使其在内容上类似于内容图像，在艺术风格上类似于风格图像。
基本原理¶
原理很简单：我们定义两个距离，一个用于内容 (\(D_C\))，一个用于风格 (\(D_S\))。 \(D_C\) 衡量两张图像在内容上的差异，而 \(D_S\) 衡量两张图像在风格上的差异。然后，我们取第三张图像，即输入图像，并对其进行变换，使其与内容图像的内容距离以及与风格图像的风格距离都最小化。现在我们可以导入必要的包并开始神经迁移。
导入包并选择设备¶
下面列出了实现神经迁移所需的包。
torch,torch.nn,numpy(使用 PyTorch 构建神经网络必不可少的包)torch.optim(高效的梯度下降)PIL,PIL.Image,matplotlib.pyplot(加载和显示图像)torchvision.transforms(将 PIL 图像转换为张量)torchvision.models(训练或加载预训练模型)copy(用于深度复制模型；系统包)
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from PIL import Image
import matplotlib.pyplot as plt
import torchvision.transforms as transforms
from torchvision.models import vgg19, VGG19_Weights
import copy
接下来，我们需要选择在哪种设备上运行网络，并导入内容和风格图像。在大型图像上运行神经迁移算法耗时较长，在 GPU 上运行时速度会快很多。我们可以使用 torch.cuda.is_available() 来检测是否有可用的 GPU。然后，我们设置 torch.device 用于整个教程。此外，.to(device) 方法用于将张量或模块移动到所需的设备。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.set_default_device(device)
加载图像¶
现在我们将导入风格图像和内容图像。原始 PIL 图像的值在 0 到 255 之间，但转换为 torch 张量后，其值会转换为 0 到 1 之间。图像还需要调整大小以具有相同的尺寸。需要注意一个重要的细节：torch 库中的神经网络使用值范围为 0 到 1 的张量进行训练。如果你尝试使用 0 到 255 的张量图像馈送网络，那么激活的特征图将无法感知预期的内容和风格。然而，Caffe 库中的预训练网络是使用 0 到 255 的张量图像进行训练的。
注意
这里是运行本教程所需图像的下载链接：picasso.jpg 和 dancing.jpg。下载这两张图像并将它们添加到当前工作目录中名为 images 的目录中。
# desired size of the output image
imsize = 512 if torch.cuda.is_available() else 128 # use small size if no GPU
loader = transforms.Compose([
transforms.Resize(imsize), # scale imported image
transforms.ToTensor()]) # transform it into a torch tensor
def image_loader(image_name):
image = Image.open(image_name)
# fake batch dimension required to fit network's input dimensions
image = loader(image).unsqueeze(0)
return image.to(device, torch.float)
style_img = image_loader("./data/images/neural-style/picasso.jpg")
content_img = image_loader("./data/images/neural-style/dancing.jpg")
assert style_img.size() == content_img.size(), \
"we need to import style and content images of the same size"
现在，让我们创建一个函数来显示图像，该函数通过将图像副本重新转换为 PIL 格式并使用 plt.imshow 显示该副本。我们将尝试显示内容图像和风格图像，以确保它们已正确导入。
unloader = transforms.ToPILImage() # reconvert into PIL image
plt.ion()
def imshow(tensor, title=None):
image = tensor.cpu().clone() # we clone the tensor to not do changes on it
image = image.squeeze(0) # remove the fake batch dimension
image = unloader(image)
plt.imshow(image)
if title is not None:
plt.title(title)
plt.pause(0.001) # pause a bit so that plots are updated
plt.figure()
imshow(style_img, title='Style Image')
plt.figure()
imshow(content_img, title='Content Image')
损失函数¶
内容损失¶
内容损失是一个函数，它代表了单个层的内容距离的加权版本。该函数接受处理输入 \(X\) 的网络中层 \(L\) 的特征图 \(F_{XL}\)，并返回图像 \(X\) 与内容图像 \(C\) 之间的加权内容距离 \(w_{CL}.D_C^L(X,C)\)。为了计算内容距离，函数必须知道内容图像的特征图 (\(F_{CL}\))。我们将此函数实现为一个 torch 模块，其构造函数接受 \(F_{CL}\) 作为输入。\(\|F_{XL} - F_{CL}\|^2\) 距离是两组特征图之间的均方误差，可以使用 nn.MSELoss 计算。
我们将此内容损失模块直接添加到用于计算内容距离的卷积层之后。这样，每次网络接收到输入图像时，内容损失都会在所需层计算，并且由于自动微分，所有梯度都将被计算。现在，为了使内容损失层透明，我们必须定义一个 forward 方法，该方法计算内容损失然后返回层的输入。计算出的损失作为模块的参数保存。
class ContentLoss(nn.Module):
def __init__(self, target,):
super(ContentLoss, self).__init__()
# we 'detach' the target content from the tree used
# to dynamically compute the gradient: this is a stated value,
# not a variable. Otherwise the forward method of the criterion
# will throw an error.
self.target = target.detach()
def forward(self, input):
self.loss = F.mse_loss(input, self.target)
return input
注意
重要细节：尽管此模块命名为 ContentLoss，但它并不是真正的 PyTorch 损失函数。如果你想将内容损失定义为 PyTorch 损失函数，你必须创建一个 PyTorch autograd 函数，以便在 backward 方法中手动重新计算/实现梯度。
风格损失¶
风格损失模块的实现方式与内容损失模块类似。它将作为网络中的一个透明层，计算该层的风格损失。为了计算风格损失，我们需要计算格拉姆矩阵 \(G_{XL}\)。格拉姆矩阵是将给定矩阵与其转置矩阵相乘的结果。在此应用中，给定矩阵是层 \(L\) 的特征图 \(F_{XL}\) 的重塑版本。\(F_{XL}\) 被重塑为 \(\hat{F}_{XL}\)，这是一个 \(K\)x\(N\) 矩阵，其中 \(K\) 是层 \(L\) 的特征图数量，\(N\) 是任意向量化特征图 \(F_{XL}^k\) 的长度。例如，\(\hat{F}_{XL}\) 的第一行对应于第一个向量化特征图 \(F_{XL}^1\)。
最后，格拉姆矩阵必须通过将其每个元素除以矩阵中的总元素数来归一化。这种归一化是为了抵消 \(\hat{F}_{XL}\) 矩阵的 \(N\) 维度较大时，格拉姆矩阵中的值也会较大的事实。这些较大的值会导致前几层（池化层之前）在梯度下降过程中产生更大的影响。风格特征往往存在于网络的更深层，因此这一归一化步骤至关重要。
def gram_matrix(input):
a, b, c, d = input.size() # a=batch size(=1)
# b=number of feature maps
# (c,d)=dimensions of a f. map (N=c*d)
features = input.view(a * b, c * d) # resize F_XL into \hat F_XL
G = torch.mm(features, features.t()) # compute the gram product
# we 'normalize' the values of the gram matrix
# by dividing by the number of element in each feature maps.
return G.div(a * b * c * d)
现在风格损失模块看起来几乎与内容损失模块完全相同。风格距离也使用 \(G_{XL}\) 和 \(G_{SL}\) 之间的均方误差计算。
class StyleLoss(nn.Module):
def __init__(self, target_feature):
super(StyleLoss, self).__init__()
self.target = gram_matrix(target_feature).detach()
def forward(self, input):
G = gram_matrix(input)
self.loss = F.mse_loss(G, self.target)
return input
导入模型¶
现在我们需要导入一个预训练的神经网络。我们将使用一个 19 层的 VGG 网络，类似于论文中使用的网络。
PyTorch 对 VGG 的实现是一个模块，分为两个子 Sequential 模块：features（包含卷积层和池化层）和 classifier（包含全连接层）。我们将使用 features 模块，因为我们需要单个卷积层的输出以测量内容和风格损失。一些层在训练期间和评估期间的行为不同，因此我们必须使用 .eval() 将网络设置为评估模式。
cnn = vgg19(weights=VGG19_Weights.DEFAULT).features.eval()
Downloading: "https://download.pytorch.org/models/vgg19-dcbb9e9d.pth" to /var/lib/ci-user/.cache/torch/hub/checkpoints/vgg19-dcbb9e9d.pth
0%| | 0.00/548M [00:00<?, ?B/s]
8%|7 | 42.1M/548M [00:00<00:01, 442MB/s]
16%|#5 | 86.5M/548M [00:00<00:01, 455MB/s]
24%|##3 | 130M/548M [00:00<00:01, 417MB/s]
32%|###1 | 174M/548M [00:00<00:00, 433MB/s]
39%|###9 | 216M/548M [00:00<00:00, 423MB/s]
47%|####6 | 257M/548M [00:00<00:00, 403MB/s]
54%|#####4 | 298M/548M [00:00<00:00, 414MB/s]
63%|######2 | 343M/548M [00:00<00:00, 429MB/s]
71%|####### | 387M/548M [00:00<00:00, 440MB/s]
78%|#######8 | 430M/548M [00:01<00:00, 396MB/s]
85%|########5 | 468M/548M [00:01<00:00, 391MB/s]
93%|#########2| 508M/548M [00:01<00:00, 398MB/s]
100%|#########9| 548M/548M [00:01<00:00, 403MB/s]
100%|##########| 548M/548M [00:01<00:00, 413MB/s]
此外，VGG 网络是在图像上训练的，每个通道通过 mean=[0.485, 0.456, 0.406] 和 std=[0.229, 0.224, 0.225] 进行归一化。我们将在将图像输入网络之前使用它们来归一化图像。
cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406])
cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225])
# create a module to normalize input image so we can easily put it in a
# ``nn.Sequential``
class Normalization(nn.Module):
def __init__(self, mean, std):
super(Normalization, self).__init__()
# .view the mean and std to make them [C x 1 x 1] so that they can
# directly work with image Tensor of shape [B x C x H x W].
# B is batch size. C is number of channels. H is height and W is width.
self.mean = torch.tensor(mean).view(-1, 1, 1)
self.std = torch.tensor(std).view(-1, 1, 1)
def forward(self, img):
# normalize ``img``
return (img - self.mean) / self.std
一个 Sequential 模块包含一个有序的子模块列表。例如，vgg19.features 包含一个按深度顺序排列的序列（Conv2d, ReLU, MaxPool2d, Conv2d, ReLU……）。我们需要在内容损失和风格损失层需要检测的卷积层之后立即添加它们。为此，我们必须创建一个新的 Sequential 模块，并在其中正确插入内容损失和风格损失模块。
# desired depth layers to compute style/content losses :
content_layers_default = ['conv_4']
style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
def get_style_model_and_losses(cnn, normalization_mean, normalization_std,
style_img, content_img,
content_layers=content_layers_default,
style_layers=style_layers_default):
# normalization module
normalization = Normalization(normalization_mean, normalization_std)
# just in order to have an iterable access to or list of content/style
# losses
content_losses = []
style_losses = []
# assuming that ``cnn`` is a ``nn.Sequential``, so we make a new ``nn.Sequential``
# to put in modules that are supposed to be activated sequentially
model = nn.Sequential(normalization)
i = 0 # increment every time we see a conv
for layer in cnn.children():
if isinstance(layer, nn.Conv2d):
i += 1
name = 'conv_{}'.format(i)
elif isinstance(layer, nn.ReLU):
name = 'relu_{}'.format(i)
# The in-place version doesn't play very nicely with the ``ContentLoss``
# and ``StyleLoss`` we insert below. So we replace with out-of-place
# ones here.
layer = nn.ReLU(inplace=False)
elif isinstance(layer, nn.MaxPool2d):
name = 'pool_{}'.format(i)
elif isinstance(layer, nn.BatchNorm2d):
name = 'bn_{}'.format(i)
else:
raise RuntimeError('Unrecognized layer: {}'.format(layer.__class__.__name__))
model.add_module(name, layer)
if name in content_layers:
# add content loss:
target = model(content_img).detach()
content_loss = ContentLoss(target)
model.add_module("content_loss_{}".format(i), content_loss)
content_losses.append(content_loss)
if name in style_layers:
# add style loss:
target_feature = model(style_img).detach()
style_loss = StyleLoss(target_feature)
model.add_module("style_loss_{}".format(i), style_loss)
style_losses.append(style_loss)
# now we trim off the layers after the last content and style losses
for i in range(len(model) - 1, -1, -1):
if isinstance(model[i], ContentLoss) or isinstance(model[i], StyleLoss):
break
model = model[:(i + 1)]
return model, style_losses, content_losses
接下来，我们选择输入图像。你可以使用内容图像的副本或白噪声。
input_img = content_img.clone()
# if you want to use white noise by using the following code:
#
# .. code-block:: python
#
# input_img = torch.randn(content_img.data.size())
# add the original input image to the figure:
plt.figure()
imshow(input_img, title='Input Image')
梯度下降¶
正如该算法的作者 Leon Gatys 在此处建议的那样，我们将使用 L-BFGS 算法来运行梯度下降。与训练网络不同，我们希望训练输入图像以最小化内容/风格损失。我们将创建一个 PyTorch L-BFGS 优化器 optim.LBFGS，并将我们的图像作为要优化的张量传递给它。
def get_input_optimizer(input_img):
# this line to show that input is a parameter that requires a gradient
optimizer = optim.LBFGS([input_img])
return optimizer
最后，我们必须定义一个执行神经迁移的函数。对于网络的每次迭代，都会馈送一个更新的输入并计算新的损失。我们将运行每个损失模块的 backward 方法来动态计算它们的梯度。优化器需要一个“闭包”函数，该函数重新评估模块并返回损失。
我们还需要解决最后一个限制。网络可能会尝试优化输入，使其值超出图像的 0 到 1 张量范围。每次运行网络时，我们可以通过将输入值校正到 0 到 1 之间来解决此问题。
def run_style_transfer(cnn, normalization_mean, normalization_std,
content_img, style_img, input_img, num_steps=300,
style_weight=1000000, content_weight=1):
"""Run the style transfer."""
print('Building the style transfer model..')
model, style_losses, content_losses = get_style_model_and_losses(cnn,
normalization_mean, normalization_std, style_img, content_img)
# We want to optimize the input and not the model parameters so we
# update all the requires_grad fields accordingly
input_img.requires_grad_(True)
# We also put the model in evaluation mode, so that specific layers
# such as dropout or batch normalization layers behave correctly.
model.eval()
model.requires_grad_(False)
optimizer = get_input_optimizer(input_img)
print('Optimizing..')
run = [0]
while run[0] <= num_steps:
def closure():
# correct the values of updated input image
with torch.no_grad():
input_img.clamp_(0, 1)
optimizer.zero_grad()
model(input_img)
style_score = 0
content_score = 0
for sl in style_losses:
style_score += sl.loss
for cl in content_losses:
content_score += cl.loss
style_score *= style_weight
content_score *= content_weight
loss = style_score + content_score
loss.backward()
run[0] += 1
if run[0] % 50 == 0:
print("run {}:".format(run))
print('Style Loss : {:4f} Content Loss: {:4f}'.format(
style_score.item(), content_score.item()))
print()
return style_score + content_score
optimizer.step(closure)
# a last correction...
with torch.no_grad():
input_img.clamp_(0, 1)
return input_img
最后，我们可以运行算法了。
output = run_style_transfer(cnn, cnn_normalization_mean, cnn_normalization_std,
content_img, style_img, input_img)
plt.figure()
imshow(output, title='Output Image')
# sphinx_gallery_thumbnail_number = 4
plt.ioff()
plt.show()
Building the style transfer model..
/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/utils/_device.py:104: UserWarning:
To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).
Optimizing..
run [50]:
Style Loss : 4.186186 Content Loss: 4.187174
run [100]:
Style Loss : 1.130905 Content Loss: 3.028577
run [150]:
Style Loss : 0.713207 Content Loss: 2.648742
run [200]:
Style Loss : 0.486411 Content Loss: 2.495289
run [250]:
Style Loss : 0.346993 Content Loss: 2.406199
run [300]:
Style Loss : 0.267672 Content Loss: 2.350107
脚本总运行时间： ( 0 分 11.353 秒)
