跳至主要内容

使用 TensorFlow 2.0 实现神经风格迁移

在本教程中，您将学习如何使用 TensorFlow 2.0 实现如 Prisma 这类强大的应用。
· 11分钟 阅读

使用 AI 探索

卷积神经网络是用于图像分类与分析最强大的神经网络体系之一。它们的应用已突破诸多边界，被证明是我们今天所见许多深度学习应用的关键要素。从高层次来看，CNN 能学习输入图像的内部特征表示，这正是其强大之处。事实证明，CNN 的这一显著特性不仅适用于图像分类等任务，也适用于图像生成。诸如 Deep Dream 和 神经风格迁移 等应用，都是基于 CNN 各层的激活及其提取的特征来合成图像的。

在本教程中，您将学习神经风格迁移的工作原理，以及如何使用 TensorFlow 2.0 来实现它。

注意：要跟进本教程，您需要了解 CNN 的工作机制。如果您需要快速回顾 CNN，推荐尝试以下资源：

什么是神经风格迁移？

对于刚入门计算机视觉与深度学习领域的人来说，神经风格迁移这个术语可能有些难以把握。要理解这个术语的各个组成部分，请看下面两张图片：

neural style transfer 1

在神经风格迁移的语境下，左边的图称为内容图，右边的图称为风格图。您的目标是用一幅图（此处为左图）去应用另一幅图（右图）的风格。这便构成了术语后两个词——风格迁移。为完成这一过程，需要训练一个神经网络（CNN）去优化一个自定义损失函数，这对应第一个词——神经。当用神经风格迁移将上述两张图片融合后，最终输出如（右图）所示：

neural style transfer 2

您可能会问，神经风格迁移有什么用？想象您有一幅文森特·梵高的画作，您希望看看如果把他的画风应用到您自己的作品上会是什么样子。这正是神经风格迁移的用武之地。另一个例子是诸如 Prisma 等多款照片滤镜应用，它们通过流畅的用户界面让您完成神经风格迁移。

内容与风格的分离

到目前为止，您对神经风格迁移的理解可能是：这一过程就是将一幅图像的内容与另一幅图像的风格进行融合。事实上，这完全正确。本节我们将理解在 CNN 的语境中，内容与风格分别意味着什么。

CNN 通常由若干卷积层与池化层组成。卷积层负责从给定图像中提取高度复杂的特征，而池化层会丢弃对图像分类问题不那么相关的细粒度空间信息。其效果是帮助 CNN 学习图像的内容，而非某些具体属性（如颜色、纹理等）。随着网络加深，特征的复杂度增加，更深的卷积层通常被视为内容表示。

风格的例子则是与图像属性相关的具体特征，如纹理、颜色等——例如在一幅画中对某种颜色的突出使用。那么，如何提取图像的风格？做法是计算卷积层之间的相关性。相关性度量两个或多个变量之间的相似/关联程度。为理解这一点，考虑一个已学习的卷积层，它由多个特征图组成。对于每个特征图，您可以度量其检测到的特征与同层其他特征图之间的相关程度。这就能估计诸如：第一张特征图中检测到的某种颜色是否与另一张特征图中的颜色相似？是否存在跨特征图共同的形状？这些特征/相似性便定义了图像的风格。在一个卷积层内比较多个特征图内容的相似性，帮助网络学习聚焦于纹理、颜色等空间特征的多尺度图像表示。

需要牢记的一点是，在应用神经风格迁移时，您还必须确保在引入所需风格的同时，原图像的内容得到保留。稍后您将看到如何实现这一点。下一节您将学习如何从图像（内容图）中提取特征并计算内容损失。

量化内容图并计算内容损失

神经风格迁移算法最早由 Gatys 等人在 2015 年的论文 A Neural Algorithm of Artistic Style 中提出。不过，本教程参考的是其后续论文 Image Style Transfer Using Convolutional Neural Networks。

根据论文 Image Style Transfer Using Convolutional Neural Networks，其采用 VGG-19 架构分别从内容图与风格图中提取内容与风格特征。为获得内容特征，使用第四个卷积块（卷积层块）的第二个卷积层。为方便起见，作者将其命名为 conv4_2。一旦得到内容特征，您需要将其与目标图像进行比较，以度量内容损失。什么是目标图像？为什么需要在此计算内容损失？我们先退一步，回答这两个问题。

要把内容与风格特征合成为单张图像，您需要从一张目标图像开始，它可以是一张空白图，也可以是内容图的拷贝。接着，为让 CNN 有效学习内容与风格特征，您需要一个自定义损失函数，通过优化它来从内容图与风格图中构建出平滑的风格化图像。这个自定义损失函数本质上是两种损失的加权组合：

  • 内容损失，确保总体内容得到保留。
  • 风格损失，控制迁移到目标图像上的风格量。

理解了上述问题后，我们回到内容损失并加以定义。内容损失定义如下——

其中，$T_c$ 表示目标图像，$C_c$ 表示内容图像。该损失函数衡量的是内容图与目标图的特征相距多远。网络将以最小化该损失为目标。下一节您将看到如何确定风格损失。

风格损失与 Gram 矩阵

为确定风格损失，论文建议从以下层中获取表示（本质是数值），并计算这些层内特征图的 Gram 矩阵。

'conv1_1'
'conv2_1'
'conv3_1'
'conv4_1'
'conv5_1'

使用多层来定义风格损失的优势在于，它能够学习图像风格的多尺度表示。Gram 矩阵帮助您判断在单个卷积层内，不同特征图之间的特征有多相似。这进一步刻画了塑造图像风格的非局部信息。在此语境下，Gram 矩阵按如下方式计算。

设某一卷积输出的维度为 8x8x16，表示它有 16 张特征图。您希望用 Gram 矩阵来度量这些特征图间特征的相似性。为此，将该卷积输出的前两个维度展平为一维向量。在本例中，展平后的一维向量包含 64 个条目。对所有特征图重复此操作，于是得到一个维度为 16x64 的矩阵（记为矩阵 A）。然后对该矩阵转置（变为 64x16）并与矩阵 A 相乘。此时的 Gram 矩阵维度为 16x16。Gram 矩阵中的某个值表示对应特征图之间的相似性。该过程的示意如下：

style loss and gram matrix

为计算目标图像与风格图像的风格损失，需在各个卷积层块处，取它们 Gram 矩阵之间的均方距离。形式化地，可定义如下：

其中，T(s,i) 是目标图像在第 i 个块处的 Gram 矩阵，S(s,i) 是风格图像在第 i 个块处的 Gram 矩阵。通过 wi，您可以为不同的卷积块设置自定义权重，以获得更细致的风格表示。最后，$a$ 是一个常数，用于衡量各块内每一层的取值。接下来我们把两种损失合在一起，定义网络在神经风格迁移过程中要优化的总损失。

定义总损失

总损失只需按自定义权重相加即可——

其中，$\alpha$ 表示内容权重，$\beta$ 表示风格权重。它们帮助您在内容与迁移到目标图像的风格量之间取得良好平衡。论文中展示了不同 $\alpha$ 与 $\beta$ 组合的效果，建议您参考。至此，您已具备实现神经风格迁移的全部要素。下一节我们直接上手实现。

神经风格迁移实战

在实现部分，您将使用 TensorFlow 2.0。该版本带来了诸多出色的新特性，是迄今最全面的更新之一。若您想了解其中的一些内容，可查看这篇文章。

TensorFlow 的最新版本（撰写本教程时）为 2.0.0-beta0。如果您尚未安装，请先按照此处的说明进行安装。我们先导入所需的包，以及内容图和风格图。

# Packages
import tensorflow as tf
from tensorflow.keras.applications.vgg19 import preprocess_input
from tensorflow.keras.models import Model
import matplotlib.pyplot as plt
import numpy as np
np.random.seed(7)

%matplotlib inline

print(tf.__version__)

2.0.0-beta0

# Load the content and style images
content = plt.imread('Content.jpeg')
style = plt.imread('Style.jpg')

# Display the images
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 10))

# Content and style images side-by-side
ax1.imshow(content)
ax1.set_title('Content Image')
ax2.imshow(style)
ax2.set_title('Style Image')
plt.show()
content and style images

下一步，您将编写一个辅助函数，把两张图片加载为数值数组（因为计算机只理解数字），并重塑其形状以适配模型。

def load_image(image):
  image = plt.imread(image)
  img = tf.image.convert_image_dtype(image, tf.float32)
  img = tf.image.resize(img, [400, 400])
  # Shape -> (batch_size, h, w, d)
  img = img[tf.newaxis, :]
  return img

# Use load_image of content and style images
content = load_image('Content.jpeg')
style = load_image('Style.jpg')

# Verify the shapes
content.shape, style.shape

(TensorShape([1, 400, 400, 3]), TensorShape([1, 400, 400, 3]))

图像已重塑。现在，您将加载一个预训练的 VGG19 模型来提取特征。由于我们只用该模型提取特征，不需要模型的分类器部分。

vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
vgg.trainable = False

Downloading data from https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg19_weights_tf_dim_ordering_tf_kernels_notop.h5
80142336/80134624 [==============================] - 5s 0us/step

作者在论文中明确了用于提取内容和风格特征的块名。您刚加载的 VGG19 使用相同的命名约定。为利用这一点，我们先打印网络中所有层的名称。

# Print the layer names for convenience
for layer in vgg.layers:
  print(layer.name)

input_1
block1_conv1
block1_conv2
block1_pool
block2_conv1
block2_conv2
block2_pool
block3_conv1
block3_conv2
block3_conv3
block3_conv4
block3_pool
block4_conv1
block4_conv2
block4_conv3
block4_conv4
block4_pool
block5_conv1
block5_conv2
block5_conv3
block5_conv4
block5_pool

而我们关心用于获取风格特征的是以下这些：

'conv1_1'
'conv2_1'
'conv3_1'
'conv4_1'
'conv5_1'

对于内容特征，您需要 conv4_2。我们将它们分别存入变量。

# Content layer
content_layers = ['block4_conv2']

# Style layer
style_layers = ['block1_conv1',
                'block2_conv1',
                'block3_conv1',
                'block4_conv1',
                'block5_conv1']


num_content_layers = len(content_layers)
num_style_layers = len(style_layers)

现在创建一个由上述指定层组成的自定义 VGG 模型。这样您就可以在图像上做前向传播，并在过程中提取所需特征。

def mini_model(layer_names, model):

  outputs = [model.get_layer(name).output for name in layer_names]

  model = Model([vgg.input], outputs)
  return model

在 TensorFlow 中定义 Gram 矩阵非常容易，您可以这样做：

# Gram matrix
def gram_matrix(tensor):
  temp = tensor
  temp = tf.squeeze(temp)
  fun = tf.reshape(temp, [temp.shape[2], temp.shape[0]*temp.shape[1]])
  result = tf.matmul(temp, temp, transpose_b=True)
  gram = tf.expand_dims(result, axis=0)
  return gram

接下来，您将使用 mini_model() 定义一个自定义模型，用于从相应图像返回内容与风格特征。

 class Custom_Style_Model(tf.keras.models.Model):
  def __init__(self, style_layers, content_layers):
    super(Custom_Style_Model, self).__init__()
    self.vgg =  mini_model(style_layers + content_layers, vgg)
    self.style_layers = style_layers
    self.content_layers = content_layers
    self.num_style_layers = len(style_layers)
    self.vgg.trainable = False

  def call(self, inputs):
    # Scale back the pixel values
    inputs = inputs*255.0
    # Preprocess them with respect to VGG19 stats
    preprocessed_input = preprocess_input(inputs)
    # Pass through the mini network
    outputs = self.vgg(preprocessed_input)
    # Segregate the style and content representations
    style_outputs, content_outputs = (outputs[:self.num_style_layers],
                                      outputs[self.num_style_layers:])

    # Calculate the gram matrix for each layer
    style_outputs = [gram_matrix(style_output)
                     for style_output in style_outputs]

    # Assign the content representation and gram matrix in
    # a layer by layer fashion in dicts
    content_dict = {content_name:value
                    for content_name, value
                    in zip(self.content_layers, content_outputs)}

    style_dict = {style_name:value
                  for style_name, value
                  in zip(self.style_layers, style_outputs)}

    return {'content':content_dict, 'style':style_dict}

现在您已定义好自定义模型，让我们把它应用到图像上，分别获取内容与风格特征——

# Note that the content and style images are loaded in
# content and style variables respectively
extractor = Custom_Style_Model(style_layers, content_layers)
style_targets = extractor(style)['style']
content_targets = extractor(content)['content']


In the paper, optimization was done using the L-BFGS algorithm, but you can use Adam also.
opt = tf.optimizers.Adam(learning_rate=0.02)

现在我们来定义整体的内容与风格权重，以及先前讨论过的各风格表示的分层权重。注意：这些都是超参数，建议您多加尝试与调优。

 # Custom weights for style and content updates
style_weight=100
content_weight=10

# Custom weights for different style layers
style_weights = {'block1_conv1': 1.,
                 'block2_conv1': 0.8,
                 'block3_conv1': 0.5,
                 'block4_conv1': 0.3,
                 'block5_conv1': 0.1}

接下来进入最关键的部分，也使神经风格迁移过程更有趣——损失函数。

# The loss function to optimize
def total_loss(outputs):
    style_outputs = outputs['style']
    content_outputs = outputs['content']
    style_loss = tf.add_n([style_weights[name]*tf.reduce_mean((style_outputs[name]-style_targets[name])**2)
                           for name in style_outputs.keys()])
    # Normalize
    style_loss *= style_weight / num_style_layers

    content_loss = tf.add_n([tf.reduce_mean((content_outputs[name]-content_targets[name])**2)
                             for name in content_outputs.keys()])
    # Normalize
    content_loss *= content_weight / num_content_layers
    loss = style_loss + content_loss
    return loss

您现在将编写另一个函数，用于：

  • 计算刚定义的损失函数的梯度。
  • 利用这些梯度更新目标图像。

借助 GradientTape，您可以利用自动微分，根据函数的计算图自动计算梯度。您还将使用 tf.function 装饰器来加速运算。更多内容请参阅此处。

@tf.function()
def train_step(image):
  with tf.GradientTape() as tape:
    # Extract the features
    outputs = extractor(image)
    # Calculate the loss
    loss = total_loss(outputs)
  # Determine the gradients of the loss function w.r.t the image pixels
  grad = tape.gradient(loss, image)
  # Update the pixels
  opt.apply_gradients([(grad, image)])
  # Clip the pixel values that fall outside the range of [0,1]
  image.assign(tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0))

在训练网络之前，仅剩的一步是定义目标图像。此处我们直接使用内容图作为目标图像。

target_image = tf.Variable(content)

一切就绪，可以开始训练网络了。

epochs = 10
steps_per_epoch = 100

step = 0
for n in range(epochs):
  for m in range(steps_per_epoch):
    step += 1
    train_step(target_image)
  plt.imshow(np.squeeze(target_image.read_value(), 0))
  plt.title("Train step: {}".format(step))
  plt.show()
WARNING: Logging before flag parsing goes to stderr.
W0617 16:21:34.491543 140709216896896 deprecation.py:323] From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/math_grad.py:1205: add_dispatch_support.<locals>.wrapper (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.where in 2.0, which has the same broadcast rule as np.where

建议您调节内容与风格权重，观察目标图像的变化。

结语

本教程参考的论文并未涉及进一步优化生成图像质量。正如您可能在上述结果中观察到的，生成图像的整体内容与风格之间在空间关系上仍需要一定的平滑。为解决这一问题，引入了 total_variantional_loss，其作用类似于正则化。总体而言，total_variational_loss 会惩罚原始神经风格迁移算法引入的高频伪影。若您有兴趣实现它，请查看本教程。以下是撰写本教程时参考的一些资源：

如果您希望提升深度学习技能，建议查看以下 DataCamp 课程与教程：

主题

Python深度学习