Neural Style Transfer with AdaIN
作者: Aritra Roy Gosthipaty, Ritwik Raha
创建日期: 2021/11/08
最后修改: 2021/11/08
描述: 使用自适应实例归一化的神经风格转移。
引言
神经风格转移 是将一幅图像的风格转移到另一幅图像内容上的过程。这一方法首次在Gatys等人的开创性论文 "A Neural Algorithm of Artistic Style"中提出。这项技术的一个主要限制是其运行时间，因为算法使用了一种缓慢的迭代优化过程。
后续论文引入了 批量归一化, 实例归一化和 条件实例归一化， 允许以新的方式进行风格转移，不再需要缓慢的迭代过程。
在这些论文之后，作者Xun Huang和Serge Belongie提出了 自适应实例归一化（AdaIN）， 这使得任意的风格转移能够实时进行。
在这个示例中，我们实现了自适应实例归一化 用于神经风格转移。我们在下面的图中展示了我们的AdaIN模型仅训练 30个周期后的输出。
您还可以通过这个 Hugging Face 演示 尝试使用自己的图像。
设置
我们首先导入必要的软件包。我们还设置了 种子以确保可重现性。全局变量是超参数， 我们可以根据需要进行更改。
import os
import numpy as np
import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
import tensorflow_datasets as tfds
from tensorflow.keras import layers
# 定义全局变量。
IMAGE_SIZE = (224, 224)
BATCH_SIZE = 64
# 为了时间限制，训练一个周期。
# 请至少使用30个周期以获得良好的结果。
EPOCHS = 1
AUTOTUNE = tf.data.AUTOTUNE
风格转移样本图库
对于神经风格转移，我们需要风格图像和内容图像。在 本示例中，我们将使用 历史上最佳艺术作品 作为我们的风格数据集以及 Pascal VOC 作为我们的内容数据集。
这与作者在原始论文中的实现有所不同，他们使用 WIKI-Art作为风格数据集，并 MSCOCO作为内容数据集。我们这样做是为了创建一个最小但可重复的示例。
从Kaggle下载数据集
历史上最佳艺术作品 数据集托管在Kaggle上，我们可以通过以下步骤轻松地在Colab中下载它：
- 请按照这里的说明 获取您的Kaggle API密钥，如果您还没有。
- 使用以下命令上传Kaggle API密钥。
from google.colab import files
files.upload()
- 使用以下命令将API密钥移动到正确的目录并下载数据集。
$ mkdir ~/.kaggle
$ cp kaggle.json ~/.kaggle/
$ chmod 600 ~/.kaggle/kaggle.json
$ kaggle datasets download ikarus777/best-artworks-of-all-time
$ unzip -qq best-artworks-of-all-time.zip
$ rm -rf images
$ mv resized artwork
$ rm best-artworks-of-all-time.zip artists.csv
tf.data 管道
在本节中，我们将构建项目的 tf.data 管道。
对于风格数据集，我们从文件夹中解码、转换和调整图像大小。对于内容图像，我们已经使用 tfds 模块提供了一个 tf.data 数据集。
在我们准备好风格和内容数据管道后，我们将两者结合在一起，获得模型将使用的数据管道。
def decode_and_resize(image_path):
"""从图像文件路径解码并调整图像大小。
Args:
image_path: 图像文件路径。
Returns:
一个调整大小的图像。
"""
image = tf.io.read_file(image_path)
image = tf.image.decode_jpeg(image, channels=3)
image = tf.image.convert_image_dtype(image, dtype="float32")
image = tf.image.resize(image, IMAGE_SIZE)
return image
def extract_image_from_voc(element):
"""从PascalVOC数据集中提取图像。
Args:
element: 数据字典。
Returns:
一个调整大小的图像。
"""
image = element["image"]
image = tf.image.convert_image_dtype(image, dtype="float32")
image = tf.image.resize(image, IMAGE_SIZE)
return image
# 获取风格图像的文件路径。
style_images = os.listdir("/content/artwork/resized")
style_images = [os.path.join("/content/artwork/resized", path) for path in style_images]
# 将风格图像分为训练、验证和测试集
total_style_images = len(style_images)
train_style = style_images[: int(0.8 * total_style_images)]
val_style = style_images[int(0.8 * total_style_images) : int(0.9 * total_style_images)]
test_style = style_images[int(0.9 * total_style_images) :]
# 构建风格和内容的tf.data数据集。
train_style_ds = (
tf.data.Dataset.from_tensor_slices(train_style)
.map(decode_and_resize, num_parallel_calls=AUTOTUNE)
.repeat()
)
train_content_ds = tfds.load("voc", split="train").map(extract_image_from_voc).repeat()
val_style_ds = (
tf.data.Dataset.from_tensor_slices(val_style)
.map(decode_and_resize, num_parallel_calls=AUTOTUNE)
.repeat()
)
val_content_ds = (
tfds.load("voc", split="validation").map(extract_image_from_voc).repeat()
)
test_style_ds = (
tf.data.Dataset.from_tensor_slices(test_style)
.map(decode_and_resize, num_parallel_calls=AUTOTUNE)
.repeat()
)
test_content_ds = (
tfds.load("voc", split="test")
.map(extract_image_from_voc, num_parallel_calls=AUTOTUNE)
.repeat()
)
# 将风格和内容数据集打包。
train_ds = (
tf.data.Dataset.zip((train_style_ds, train_content_ds))
.shuffle(BATCH_SIZE * 2)
.batch(BATCH_SIZE)
.prefetch(AUTOTUNE)
)
val_ds = (
tf.data.Dataset.zip((val_style_ds, val_content_ds))
.shuffle(BATCH_SIZE * 2)
.batch(BATCH_SIZE)
.prefetch(AUTOTUNE)
)
test_ds = (
tf.data.Dataset.zip((test_style_ds, test_content_ds))
.shuffle(BATCH_SIZE * 2)
.batch(BATCH_SIZE)
.prefetch(AUTOTUNE)
)
下载并准备数据集 voc/2007/4.0.0 (下载: 868.85 MiB, 生成: 未知大小, 总计: 868.85 MiB) 到 /root/tensorflow_datasets/voc/2007/4.0.0...
下载完成...: 0 url [00:00, ? url/s]
下载大小...: 0 MiB [00:00, ? MiB/s]
提取完成...: 0 文件 [00:00, ? 文件/s]
0 个示例 [00:00, ? 示例/s]
正在将示例打乱并写入 /root/tensorflow_datasets/voc/2007/4.0.0.incompleteP16YU5/voc-test.tfrecord
0%| | 0/4952 [00:00<?, ? 示例/s]
0 个示例 [00:00, ? 示例/s]
正在将示例打乱并写入 /root/tensorflow_datasets/voc/2007/4.0.0.incompleteP16YU5/voc-train.tfrecord
0%| | 0/2501 [00:00<?, ? 示例/s]
0 个示例 [00:00, ? 示例/s]
正在将示例打乱并写入 /root/tensorflow_datasets/voc/2007/4.0.0.incompleteP16YU5/voc-validation.tfrecord
0%| | 0/2510 [00:00<?, ? 示例/s]
数据集 voc 已下载并准备好到 /root/tensorflow_datasets/voc/2007/4.0.0。后续调用将重用此数据。
可视化数据
在训练之前可视化数据总是更好。为了确保我们的预处理管道的正确性，我们从数据集中可视化 10 个样本。
style, content = next(iter(train_ds))
fig, axes = plt.subplots(nrows=10, ncols=2, figsize=(5, 30))
[ax.axis("off") for ax in np.ravel(axes)]
for (axis, style_image, content_image) in zip(axes, style[0:10], content[0:10]):
(ax_style, ax_content) = axis
ax_style.imshow(style_image)
ax_style.set_title("风格图像")
ax_content.imshow(content_image)
ax_content.set_title("内容图像")
架构
风格迁移网络以内容图像和风格图像作为输入，并输出风格迁移图像。AdaIN 的作者提出了一种简单的编码器-解码器结构来实现这一点。
内容图像（C）和风格图像（S）都被输入到编码器网络中。这些编码器网络的输出（特征图）随后输入 AdaIN 层。AdaIN 层计算一个组合特征图。该特征图然后输入一个随机初始化的解码器网络，这个网络作为神经风格迁移图像的生成器。
风格特征图（fs）和内容特征图（fc）被输入到 AdaIN 层。这一层生成了组合特征图 t。函数 g 表示解码器（生成器）网络。
编码器
编码器是经过预训练的（在 imagenet 上预训练）VGG19 模型的一部分。我们从 block4-conv1 层切割模型。输出层如作者在其论文中所建议的。
def get_encoder():
vgg19 = keras.applications.VGG19(
include_top=False,
weights="imagenet",
input_shape=(*IMAGE_SIZE, 3),
)
vgg19.trainable = False
mini_vgg19 = keras.Model(vgg19.input, vgg19.get_layer("block4_conv1").output)
inputs = layers.Input([*IMAGE_SIZE, 3])
mini_vgg19_out = mini_vgg19(inputs)
return keras.Model(inputs, mini_vgg19_out, name="mini_vgg19")
自适应实例归一化
AdaIN 层接收内容和风格图像的特征。该层可以通过以下方程定义：
其中 sigma 是标准偏差，mu 是相关变量的均值。在上面的方程中，内容特征图 fc 的均值和方差与风格特征图 fs 的均值和方差对齐。
需要注意的是，作者提出的 AdaIN 层除了均值和方差之外不使用其他参数。该层也没有任何可训练参数。这就是我们使用 Python 函数 而不是使用 Keras 层 的原因。该函数接收风格和内容特征图，计算图像的均值和标准偏差，并返回自适应实例归一化特征图。
def get_mean_std(x, epsilon=1e-5):
axes = [1, 2]
# 计算张量的均值和标准偏差。
mean, variance = tf.nn.moments(x, axes=axes, keepdims=True)
standard_deviation = tf.sqrt(variance + epsilon)
return mean, standard_deviation
def ada_in(style, content):
"""计算 AdaIn 特征图。
参数:
style: 风格特征图。
content: 内容特征图。
返回:
AdaIN 特征图。
"""
content_mean, content_std = get_mean_std(content)
style_mean, style_std = get_mean_std(style)
t = style_std * (content - content_mean) / content_std + style_mean
return t
解码器
作者指明，解码器网络必须镜像编码器
network. 我们对编码器进行了对称反转以构建解码器。我们使用了 UpSampling2D 层来增加特征图的空间分辨率。
请注意，作者警告不要在解码器网络中使用任何归一化层，并确实展示了包含批归一化或实例归一化会降低整体网络的性能。
这是整个架构中唯一可以训练的部分。
def get_decoder():
config = {"kernel_size": 3, "strides": 1, "padding": "same", "activation": "relu"}
decoder = keras.Sequential(
[
layers.InputLayer((None, None, 512)),
layers.Conv2D(filters=512, **config),
layers.UpSampling2D(),
layers.Conv2D(filters=256, **config),
layers.Conv2D(filters=256, **config),
layers.Conv2D(filters=256, **config),
layers.Conv2D(filters=256, **config),
layers.UpSampling2D(),
layers.Conv2D(filters=128, **config),
layers.Conv2D(filters=128, **config),
layers.UpSampling2D(),
layers.Conv2D(filters=64, **config),
layers.Conv2D(
filters=3,
kernel_size=3,
strides=1,
padding="same",
activation="sigmoid",
),
]
)
return decoder
损失函数
在这里我们构建神经风格迁移模型的损失函数。作者建议使用预训练的 VGG-19 来计算网络的损失函数。重要的是要记住，这将用于仅训练解码器网络。总损失 (Lt) 是内容损失 (Lc) 和风格损失 (Ls) 的加权组合。lambda 项用于变化风格转移的量。
内容损失
这是内容图像特征与神经风格转移图像特征之间的欧几里得距离。
在这里，作者建议使用 AdaIn 层 t 的输出作为内容目标，而不是使用原始图像的特征作为目标。这是为了加速收敛。
风格损失
作者建议计算统计特征（均值和方差）之间的差异，而不是使用更常用的 Gram 矩阵，这样在概念上更清晰。可以通过以下方程轻松可视化：
其中 theta 表示用于计算损失的 VGG-19 层。在这种情况下，对应于：
block1_conv1block1_conv2block1_conv3block1_conv4
def get_loss_net():
vgg19 = keras.applications.VGG19(
include_top=False, weights="imagenet", input_shape=(*IMAGE_SIZE, 3)
)
vgg19.trainable = False
layer_names = ["block1_conv1", "block2_conv1", "block3_conv1", "block4_conv1"]
outputs = [vgg19.get_layer(name).output for name in layer_names]
mini_vgg19 = keras.Model(vgg19.input, outputs)
inputs = layers.Input([*IMAGE_SIZE, 3])
mini_vgg19_out = mini_vgg19(inputs)
return keras.Model(inputs, mini_vgg19_out, name="loss_net")
神经风格迁移
这是训练模块。我们将编码器和解码器封装在 tf.keras.Model 子类中。这使我们能够定制在 model.fit() 循环中发生的事情。
class NeuralStyleTransfer(tf.keras.Model):
def __init__(self, encoder, decoder, loss_net, style_weight, **kwargs):
super().__init__(**kwargs)
self.encoder = encoder
self.decoder = decoder
self.loss_net = loss_net
self.style_weight = style_weight
def compile(self, optimizer, loss_fn):
super().compile()
self.optimizer = optimizer
self.loss_fn = loss_fn
self.style_loss_tracker = keras.metrics.Mean(name="style_loss")
self.content_loss_tracker = keras.metrics.Mean(name="content_loss")
self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
def train_step(self, inputs):
style, content = inputs
# 初始化内容和风格损失。
loss_content = 0.0
loss_style = 0.0
with tf.GradientTape() as tape:
# 编码风格和内容图像。
style_encoded = self.encoder(style)
content_encoded = self.encoder(content)
# 计算 AdaIN 目标特征图。
t = ada_in(style=style_encoded, content=content_encoded)
# 生成神经风格转移图像。
reconstructed_image = self.decoder(t)
# 计算损失。
reconstructed_vgg_features = self.loss_net(reconstructed_image)
style_vgg_features = self.loss_net(style)
loss_content = self.loss_fn(t, reconstructed_vgg_features[-1])
for inp, out in zip(style_vgg_features, reconstructed_vgg_features):
mean_inp, std_inp = get_mean_std(inp)
mean_out, std_out = get_mean_std(out)
loss_style += self.loss_fn(mean_inp, mean_out) + self.loss_fn(
std_inp, std_out
)
loss_style = self.style_weight * loss_style
total_loss = loss_content + loss_style
# 计算梯度并优化解码器。
trainable_vars = self.decoder.trainable_variables
gradients = tape.gradient(total_loss, trainable_vars)
self.optimizer.apply_gradients(zip(gradients, trainable_vars))
# 更新跟踪器。
self.style_loss_tracker.update_state(loss_style)
self.content_loss_tracker.update_state(loss_content)
self.total_loss_tracker.update_state(total_loss)
return {
"style_loss": self.style_loss_tracker.result(),
"content_loss": self.content_loss_tracker.result(),
"total_loss": self.total_loss_tracker.result(),
}
def test_step(self, inputs):
style, content = inputs
# 初始化内容和风格损失。
loss_content = 0.0
loss_style = 0.0
# 编码风格和内容图像。
style_encoded = self.encoder(style)
content_encoded = self.encoder(content)
# 计算 AdaIN 目标特征图。
t = ada_in(style=style_encoded, content=content_encoded)
# 生成神经风格转移图像。
reconstructed_image = self.decoder(t)
# 计算损失。
recons_vgg_features = self.loss_net(reconstructed_image)
style_vgg_features = self.loss_net(style)
loss_content = self.loss_fn(t, recons_vgg_features[-1])
for inp, out in zip(style_vgg_features, recons_vgg_features):
mean_inp, std_inp = get_mean_std(inp)
mean_out, std_out = get_mean_std(out)
loss_style += self.loss_fn(mean_inp, mean_out) + self.loss_fn(
std_inp, std_out
)
loss_style = self.style_weight * loss_style
total_loss = loss_content + loss_style
# 更新跟踪器。
self.style_loss_tracker.update_state(loss_style)
self.content_loss_tracker.update_state(loss_content)
self.total_loss_tracker.update_state(total_loss)
return {
"style_loss": self.style_loss_tracker.result(),
"content_loss": self.content_loss_tracker.result(),
"total_loss": self.total_loss_tracker.result(),
}
@property
def metrics(self):
return [
self.style_loss_tracker,
self.content_loss_tracker,
self.total_loss_tracker,
]
训练监控回调
这个回调用于可视化模型在每个训练周期结束时的风格迁移输出。风格迁移的目标无法被很好地量化，需要通过观众的主观评估来进行。因此，可视化是评估模型的一个关键方面。
test_style, test_content = next(iter(test_ds))
class TrainMonitor(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
# 编码风格和内容图像。
test_style_encoded = self.model.encoder(test_style)
test_content_encoded = self.model.encoder(test_content)
# 计算 AdaIN 特征。
test_t = ada_in(style=test_style_encoded, content=test_content_encoded)
test_reconstructed_image = self.model.decoder(test_t)
# 绘制风格、内容和 NST 图像。
fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(20, 5))
ax[0].imshow(tf.keras.utils.array_to_img(test_style[0]))
ax[0].set_title(f"风格: {epoch:03d}")
ax[1].imshow(tf.keras.utils.array_to_img(test_content[0]))
ax[1].set_title(f"内容: {epoch:03d}")
ax[2].imshow(
tf.keras.utils.array_to_img(test_reconstructed_image[0])
)
ax[2].set_title(f"NST: {epoch:03d}")
plt.show()
plt.close()
训练模型
在这一部分，我们定义优化器、损失函数和训练模块。我们将训练模块与优化器和损失函数编译，然后进行训练。
注意: 由于时间限制，我们只训练模型一个周期，但我们需要训练至少 30 个周期以获得良好的结果。
optimizer = keras.optimizers.Adam(learning_rate=1e-5)
loss_fn = keras.losses.MeanSquaredError()
encoder = get_encoder()
loss_net = get_loss_net()
decoder = get_decoder()
model = NeuralStyleTransfer(
encoder=encoder, decoder=decoder, loss_net=loss_net, style_weight=4.0
)
model.compile(optimizer=optimizer, loss_fn=loss_fn)
history = model.fit(
train_ds,
epochs=EPOCHS,
steps_per_epoch=50,
validation_data=val_ds,
validation_steps=50,
callbacks=[TrainMonitor()],
)
从 https://storage.googleapis.com/tensorflow/keras-applications/vgg19/vgg19_weights_tf_dim_ordering_tf_kernels_notop.h5 下载数据
80142336/80134624 [==============================] - 1s 0us/step
80150528/80134624 [==============================] - 1s 0us/step
50/50 [==============================] - ETA: 0s - style_loss: 213.1439 - content_loss: 141.1564 - total_loss: 354.3002
50/50 [==============================] - 124s 2s/step - style_loss: 213.1439 - content_loss: 141.1564 - total_loss: 354.3002 - val_style_loss: 167.0819 - val_content_loss: 129.0497 - val_total_loss: 296.1316
推断
在我们训练好模型后，现在需要进行推断。我们将从测试数据集中传递任意内容和风格图像，并查看输出图像。
注意: 要尝试在您自己的图像上使用此模型，您可以使用这个 Hugging Face demo。
for style, content in test_ds.take(1):
style_encoded = model.encoder(style)
content_encoded = model.encoder(content)
t = ada_in(style=style_encoded, content=content_encoded)
reconstructed_image = model.decoder(t)
fig, axes = plt.subplots(nrows=10, ncols=3, figsize=(10, 30))
[ax.axis("off") for ax in np.ravel(axes)]
for axis, style_image, content_image, reconstructed_image in zip(
axes, style[0:10], content[0:10], reconstructed_image[0:10]
):
(ax_style, ax_content, ax_reconstructed) = axis
ax_style.imshow(style_image)
ax_style.set_title("风格图像")
ax_content.imshow(content_image)
ax_content.set_title("内容图像")
ax_reconstructed.imshow(reconstructed_image)
ax_reconstructed.set_title("NST 图像")
结论
自适应实例归一化允许实时任意风格转移。还需要注意的是，作者的新提议是通过对齐风格和内容图像的统计特征（均值和标准差）来实现这一点。
注意: AdaIN 还作为 Style-GANs 的基础。
参考
致谢
我们感谢 Luke Wood 的详细评审。