用自注意力GAN为百年旧照上色：效果惊艳，多图预警！

项目地址：https://github.com/jantic/DeOldify/blob/master/README.md

该项目的目的是为旧照片着色并将其修复。本文将简要介绍技术细节。首先，先来看一下 DeOldify 对旧照片的修复效果！（大部分原图像来自 r/TheWayWeWere subreddit。）

和人体骨骼一起拍照的医学生（约 1890 年）

Whirling Horse（1898）

19 世纪 80 年代的巴黎

20 世纪 20 年代的爱丁堡俯瞰

伦敦滑铁卢车站人们第一次看到电视机（1936）

抽鸦片的大清子民（1880）

非常旧和质量差的照片也能修复得很好：

南达科他州枯木镇（1877）

家庭合照（1877）

但，模型效果并不总是完美的。下图中红手让人抓狂：

塞内卡原住民（1908）

该模型还能对黑白线稿进行着色：

技术细节

这是一个基于深度学习的模型。具体来说，我所做的是将以下方法组合在一起：

• 带自注意力机制的生成对抗网络。生成器是一个预训练 Unet，我将它修改为具有光谱归一化和自注意力。这是一个非常简单的转换过程。不过，之前我拼命地尝试用 Wasserstein GAN，但效果并不好，直到用上这个版本，一切都变了。我喜欢 Wasserstein GAN 的理念，但它在实践中并没有成功。我爱上了带自注意力的 GAN。

• 受 GAN Progressive Growing 的启发（参见《Progressive Growing of GANs for Improved Quality, Stability, and Variation》）设计出的训练结构。不同之处在于层数保持不变——我只是不断改变输入的尺寸并调整学习率，以确保尺寸之间的转换顺利进行。似乎基本最终结果是相同的——训练更快、更稳定，且泛化效果更好。

• 两个时间尺度上的更新规则（参见《GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium》）。这个也非常简单，只是一个一对一生成器/critic 迭代及较高级别的 critic 学习率。

• 生成器损失分为两个部分：一部分是基于 VGG16 的基本感知损失（或特征损失）——这只是偏向生成器模型来复制输入图像。另一部分是来自 critic 的损失分数。感知损失本身不足以产生好的结果，只会带来大量的棕色/绿色/蓝色填充——测试的时候作弊可是神经网络的强项！这里要理解的重点是：GAN 本质上是在为你学习损失函数——这实际上是朝着我们在机器学习中追求的理想迈进了一大步。当然，让机器学习你以前手工编码的东西通常会得到更好的结果。在这个例子中就是这样。

该模型的美妙之处在于，它可能适用于所有类型的图像修复，效果也应该很好。上文展示了很多模型结果，但这只是我想开发的流程的一个组件。

接下来，我想做的是「defade」模型。我已经做了一些初步努力，截至本文写作时，它还在训练阶段。大体上，是训练同样的模型用于重建使用过分的对比度／亮度调整后的图像。我已经看到了一些还不错的试验结果：

关于该项目

该项目的重点是：我想利用 GAN 使旧照片变得好看，更重要的是，做一个有用的项目。我对视频很感兴趣，但是首先我需要先解决如何处理模型和内存的关系。如果该模型在 1080TI 上的训练时间不用花费两三天就太棒了（然而 GAN 通常需要这么长时间……）我将积极更新和改进代码，并尝试使该项目对用户友好。

开始操作

该项目基于 Fast.AI 库构建。不过，我使用的是旧版本库，即将更新为新版本。构建该项目的先决条件是：

• 旧版本 Fast.AI 库。埋头该项目两个月后，我有点疑惑，因为标注为「old」的旧版本 Fast.AI 库和我用的库不一样。因此建议使用这里的库：https://github.com/jantic/fastai。

• 不管 Fast.AI 的依赖项是什么，总有方便的 requirements.txt 和 environment.yml。

• Pytorch 0.4.1（需要 spectral_norm，因此需要最新的稳定版本）。

• Jupyter Lab

• Tensorboard（即安装 Tensorflow）和 TensorboardX (https://github.com/lanpa/tensorboardX)。

• ImageNet：很棒的训练数据集。

• 强大的 GPU：我想要比 GeForce 1080TI (11GB) 更大内存的 GPU。内存越少花费时间就会越多。Unet 和 Critic 非常大，但是它们规模越大，获得的结果也会越好。

至于想要立刻开始转换自己的图像的人：如果要立刻开始使用自己的图像而不训练模型的话，你需要我上传预训练权重。我正在做这件事。做好之后，你可以在可视化 notebook 中看到。我将使用 ColorizationVisualization.ipynb。你需要设置 colorizer_path = IMAGENET.parent/('bwc_rc_gen_192.h5')。我上传的权重文件用于生成器（着色器）。

然后你就可以把想转换的任意图像放在/test_images/文件夹，然后在 notebook 内部看到结果：

vis.plot_transformed_image("test_images/derp.jpg", netG, md.val_ds, tfms=x_tfms, sz=500)

我把图像大小设置为 500px 左右，你需要在足够内存的 GPU 上运行（比如 11 GB GeForce 1080Ti）。如果内存过少，你需要将图像调小或者尝试在 CPU 上运行模型。我试过后一种方法，非常慢……我还没有研究具体原因。