我们在 Facebook f8 大会上展示了这项工作。您可以在这里观看我们演讲的视频,或者继续阅读以了解更多细节和示例。
去“渣”、去“旧”和超分辨率
在本文中,我们将介绍“去渣”(decrappification)的概念,这是一种在fastai上基于PyTorch实现的深度学习方法,可以做一些非常了不起的事情,比如... 给经典黑白电影上色——甚至包括默片时代的电影,就像下面这个例子
同样的方法可以使您的旧家庭照片看起来像是在现代相机上拍摄的,甚至可以提高在索尔克研究所(Salk Institute)使用最先进设备拍摄的显微镜图像的清晰度,从而使细胞分析的准确性提高300%。
“去渣”(Decrappify)的起源
生成模型是生成音乐、图像、文本及其他复杂数据类型的模型。近年来,生成模型的进步令人震惊,这很大程度上归功于深度学习,特别是生成对抗模型(GANs)。然而,GANs 众所周知地难以训练,原因在于需要大量数据、需要许多 GPU 和大量训练时间,并且对微小的超参数变化高度敏感。
fast.ai 近年来一直致力于使一系列模型更容易、更快地训练,特别注重使用迁移学习。迁移学习指的是使用现成数据和快速易于计算的损失函数对模型进行预训练,然后针对可能具有较少标签或计算成本较高的任务对该模型进行微调。这似乎是 GAN 训练问题的一个潜在解决方案,因此在2018年底,fast.ai 开发了一种用于生成建模的迁移学习技术。
fast.ai 选择的预训练模型是这样的:从一个图像数据集开始,“加渣”(crappify)图像,例如降低分辨率、添加 jpeg 失真、用随机文本遮挡部分内容。然后训练一个模型来“去渣”(decrappify)这些图像,使它们恢复到原始状态。fast.ai 从一个经过 ImageNet 分类预训练的模型开始,并添加了一个U-Net 上采样网络,并对常规 U-Net 添加了各种现代调整。最初使用了简单的快速损失函数:像素的均方误差。这个 U-Net 可以在短短几分钟内训练完成。然后,损失函数被替换为生成建模文献中使用的其他损失函数的组合(f8 视频中有更多细节),并又训练了几个小时。当时的计划是最后再加入 GAN 进行最后的几个 epoch——然而结果证明,结果非常好,以至于 fast.ai 最终没有在最终模型中使用 GAN。
“去旧”(DeOldify)的起源
DeOldify 与 fast.ai 开始研究 decrappification 大约同时开发,旨在为黑白照片上色。Jason Antic 观看了2018年春季 fast.ai 课程,该课程介绍了 GANs、U-Nets 和其他技术,并思考将它们结合用于上色会发生什么。Jason 最初使用 GANs 的实验大部分都失败了,所以他尝试了别的东西——自注意力 GAN(SAGAN)。他的目标是能够成功地为真实世界的老旧图像上色,处理胶片退化引起的噪声、对比度和亮度问题。模型需要用模拟这些问题的照片进行训练。为此,他从 ImageNet 数据集中的图像开始,将它们转换为黑白,然后添加随机的对比度、亮度和其他变化。换句话说,他也在“加渣”图像!
结果令人惊叹,互联网上到处都在讨论 Jason 新的“DeOldify”程序。Jeremy 看到了一些早期结果,很高兴看到有人在图像生成方面取得了很好的结果。他联系了 Jason 以了解更多信息。Jeremy 和 Jason 很快意识到他们都在使用非常相似的技术,但也都朝着一些不同的方向发展。于是他们决定联手开发一个包含他们所有最佳想法的 decrappification 流程。
联手的结果是一个完全跳过 GANs 的过程,可以在游戏 PC 上进行训练。Jason 的所有开发都是在餐厅里的 Linux 主机上完成的,每次实验只使用一个消费级 GPU(一块 GeForce 1080Ti)。缺乏令人印象深刻的硬件和工业资源并没有阻止切实的进步。事实上,这可能反而促进了它。
Jason 随后将这个过程进一步推进,从静止图像转向电影。他发现,在与 fast.ai 开发的过程基础上进行微量的 GAN 微调,就可以在短短几个小时内创建彩色电影,其质量超越了之前任何自动化过程所能达到的水平。
显微镜超分辨率的起源
与此同时,索尔克研究所 Waitt 先进生物光子学核心(WABC)主任 Uri Manor 正在寻找方法,以同时提高 WABC 最先进的 ZEISS 扫描电子显微镜和激光扫描共聚焦显微镜拍摄的图像的分辨率、速度和信噪比。这三个参数之间明显存在矛盾——这就是所谓的“妥协三角形”的一种变体,也是所有摄影师和成像科学家存在的痛点。WABC 的先进显微镜被索尔克(以及包括 Scripps 和 UCSD 在内的几个邻近机构)的研究人员广泛使用,用于研究生命体的超微结构组织和动力学,范围从植物组织中的碳捕获机器到大脑回路中的突触连接,再到癌细胞和神经元中的能量生成线粒体。扫描电子显微镜的特点是能够连续切片和成像整个组织块,从而以接近纳米分辨率获得三维体积数据集。WABC 的所谓“Airyscan”扫描共聚焦显微镜拥有尖端的32个六边形排列的探测器阵列,可以实现接近普通显微镜两倍分辨率的荧光成像,同时还提供8倍的灵敏度和速度。
得益于 Wicklow AI 医学研究计划(WAMRI),Jeremy Howard 和 Fred Monroe 得以访问索尔克,参观那里的一些杰出工作,并讨论利用深度学习帮助索尔克一些项目的机会。见到 Uri 后,立即清楚 fast.ai 的技术非常适合 Uri 对高分辨率显微镜的需求。Fred、Uri 以及一个由索尔克领导的、来自 UCSD 到 UT-Austin 的科学家团队合作,将这些方法引入显微镜领域,结果令人惊叹。通过仔细获取用于训练的高分辨率图像,该小组验证了用于电子显微镜和荧光显微镜图像超分辨率处理的“通用”模型,实现了比以往报道的更快的成像速度、更高的吞吐量、更低的样本损伤和更小的文件大小。由于这些模型能够恢复在相对低成本显微镜上获取的图像,因此这个模型也为那些无法负担最新尖端仪器的精英机构之外的人们提供了“民主化”高分辨率成像的机会。
对于创建显微镜视频,Fred 使用了与 Jason 用于经典好莱坞电影不同的方法。他受到这篇关于稳定视频中的神经风格迁移的博客文章的启发,能够在用于创建单图像超分辨率的损失函数中添加一个“稳定性”度量。这种稳定性损失度量鼓励图像在少量随机噪声存在的情况下保持稳定的特征。噪声注入本来就是索尔克创建训练集过程的一部分——所以这是一个简单的修改。这种稳定性与有关视频前一帧和后一帧的信息相结合,显著减少了闪烁,并提高了处理低分辨率视频时的输出质量。有关该过程的更多细节,请参见下面的章节——“创建超分辨率显微镜视频注意事项”。
深入了解 DeOldify
让我们详细看看DeOldify幕后发生的事情。但首先,您可以自己使用 DeOldify!最简单的方法是使用这些免费的 Colab 笔记本,它们将引导您完成整个过程
或者,您可以从GitHub 仓库下载代码并在本地运行。
最先进技术的进展
Zhang 等人的“Colorful Image Colorization”模型目前很流行,被广泛使用,并且之前一直是最先进的技术。下面展示了原始黑白照片(左),以及“Colorful Image Colorization”模型(中)与最新版本DeOldify(右)的对比。请注意,DeOldify 照片中的人物和物体上色更一致、准确且细节更丰富。DeOldify 生成的图像通常可以被认为是接近照片级的真实感。
此外,DeOldify 中的高质量视频上色之所以成为可能,得益于训练方面的进展,这些进展大大提高了渲染的稳定性和质量,这主要归功于采用了 NoGAN 训练。以下片段展示了 DeOldify 不仅能很好地为视频(甚至特效!)上色,还能保持帧间的时序一致性。
DeOldify 的设计
DeOldify 中有一些关键的设计决策对渲染图像的质量产生了显著影响。
自注意力
DeOldify 最重要的设计选择之一是使用自注意力,这在“自注意力生成对抗网络(Self-Attention Generative Adversarial Networks)”(SAGAN)论文中实现。该论文总结了使用它们的动机
“传统的卷积 GANs 生成高分辨率细节的功能仅依赖于低分辨率特征图中的空间局部点。在 SAGAN 中,可以使用所有特征位置的线索生成细节。此外,判别器可以检查图像远处部分的高细节特征是否相互一致。”
同样的这种方法被应用于 DeOldify 中使用的判别器和基于U-Net的生成器。其动机很简单:你希望在上色时达到最大的连续性、一致性和完整性,而自注意力对此至关重要。在没有自注意力的模型中,当你试图为包含大特征(如海水)的图像上色时,这个问题尤为明显。在这种模型中,你经常会看到这些大面积区域填充了不一致的颜色。
Zhang 等人的“Colorful Image Colorization”模型输出在左边,DeOldify 输出在右边。请注意,背景中的水以及左边男子的衣物,在 DeOldify 模型中上色得更加一致、完整和准确。这在很大程度上归功于自注意力,通过自注意力,上色决策可以轻松地考虑到更全局的特征背景。
自注意力还有助于在着色中实现极高的细节水平。
可靠的特征检测
与其他上色模型不同,DeOldify 为其每个生成器模型使用带有预训练 resnet 主干网络的自定义 U-Nets。这些基于 Fast.AI 精心设计的 DynamicUnet,并进行了一些微小修改。与标准 U-Nets 的不同之处包括前面提到的自注意力,以及添加了谱归一化。这些更改是参照“自注意力生成对抗网络”论文中完成的工作建模的。
“视频”和“稳定”模型使用 resnet101 主干网络,解码器端侧重于宽度(滤波器数量)而非深度(层数)。这种配置已被证明能支持目前看到的最稳定和可靠的渲染效果。相比之下,“艺术”模型使用 resnet34 主干网络,解码器端侧重于深度而非宽度。这种配置非常适合创建有趣的着色和高细节渲染,但代价是渲染的一致性不如“稳定”和“视频”模型。
使用预训练 U-Net 的两个主要潜在动机。首先,它节省了不必要的训练时间,因为上色任务中一项重要的子任务——目标识别——已经可以免费获得训练。这是基于 ImageNet 的目标识别,对于单个 GPU 来说,从头开始训练
为了进一步增强处理旧图像和低质量图像及胶片的稳健性,我们使用相当极端的亮度和对比度增强进行训练。我们还在视频模型训练中采用了高斯噪声增强,以减少模型对胶片中无关噪声(颗粒)的敏感性。
当特征识别失败时,可能会出现令人不快的渲染失败,例如“僵尸手”。
NoGAN 训练
NoGAN 是我们开发的一种新的、令人兴奋的 GAN 训练技术,旨在追求更高质量和更稳定的渲染。它如何工作,以及
以下是 NoGAN 的训练过程
预训练生成器。 生成器首先以一种更传统、更容易控制的方式进行训练——只使用感知损失(也称特征损失)。GAN 训练此时尚未引入。在这一点上,您正以最简单的方式尽最大努力训练生成器。这占据了 NoGAN 训练的大部分时间。记住:仅凭这个预训练就能让生成器模型走得很远。作为一项任务,着色将得到很好的训练,尽管颜色会偏向暗淡的色调。自注意力在这个阶段也会得到很好的训练,这非常重要。保存来自预训练生成器生成的图像。
将判别器预训练为二分类器。 就像预训练生成器一样,我们在这一步的目标是以更“传统”、更容易控制的方式尽可能多地训练判别器。没有什么比二分类器更容易了!在这里,我们将判别器训练成真实图像和伪造图像的二分类器,其中伪造图像是上一步中保存的图像。这里需要记住的一点是,您只需使用已预训练好的、用于其他图像到图像任务的判别器并对其进行微调即可。超分辨率就已这么做了,判别器的预训练权重是从用于着色任务训练的判别器那里加载的。在这种情况下,要利用预训练判别器,只需进行少量微调即可。在(几乎)正常的 GAN 设置中训练生成器和判别器。快速! 这就是令人惊讶的部分。事实证明,在这种预训练场景下,判别器在 GAN 训练期间会快速地驱动生成器进行调整。这发生在一个“拐点”到来之前的短暂窗口期。过了这个点之后,再以这种方式继续训练似乎几乎没有任何益处。事实上,如果过了这个点还继续训练,你就会开始在渲染中看到伪影和故障的出现。
在 DeOldify 的情况下,训练到这一点只需要迭代 ImageNet 数据的约 1% 到 3%(或在批量大小为五的情况下进行约 2600 到 7800 次迭代)。这仅相当于大约 30-90 分钟的 GAN 训练,与之前进行的、渐进式大小的 GAN 训练需要的三到五天形成鲜明对比。令人惊讶的是,在这短短的训练时间内,渲染质量的变化是巨大的。事实上,这就是视频模型 GAN 训练的全部内容。“艺术”和“稳定”模型更进一步,重复 NoGAN 训练过程的步骤 2-4,直到没有明显的益处(大约重复五次)。
请注意,1.4% 的帧被认为是“拐点”之前的时间点——在此之后,开始出现伪影和不正确的着色。在这种情况下,演员的皮肤变得越来越橙色和过饱和,这是不理想的。这些是使用 1e-5 的学习率生成的。当前 DeOldify 的视频模型以 5e-6 的学习率进行训练,以便更容易找到“拐点”。
关于 NoGAN 训练的研究仍在进行中,因此仍有不少问题需要深入探讨。首先,这项技术似乎能很好地适应小批量大小。视频模型是使用批量大小为 5 进行训练的(模型使用了批量归一化)。然而,“拐点”后出现的伪影和变色问题仍然存在,并且怀疑可以通过减轻批量大小相关问题来减少或消除。这可以通过增加批量大小,或者可能通过使用不受批量大小影响的归一化来实现。很可能一旦这些问题得到解决,就可以在回报递减点到来之前完成额外的有效训练。
NoGAN 的另一个未解决问题是该方法的适用范围有多广。可以推断,这应该适用于大多数图像到图像的任务,甚至可能适用于非图像相关的训练。然而,这一点还没有得到足够的探索,无法得出强有力的结论。我们在将 NoGAN 应用于超分辨率方面确实取得了有趣且令人印象深刻的结果。仅仅经过十五分钟的直接 GAN 训练(或 1500 次迭代),基于 Fast.AI 第一部分第七课特征损失的超分辨率训练输出明显变得更清晰(原始课程 Jupyter notebook 在此)。
最后,NoGAN 训练的最佳实践尚未完全探索。值得再次提及,“艺术”和“稳定”模型不仅训练了一次,而是重复了多次 NoGAN 循环。目前尚不清楚能够获得收益的重复次数,以及如何通过自动检测“拐点”之前的早期停止点来减少训练过程的繁琐。目前,拐点的确定是一个手动过程,需要人工在模型检查点处目视评估生成的图像。这些检查点需要至少每总数据的 0.1% 保存一次,否则很容易错过。这确实很繁琐,而且容易出错。
如何实现稳定视频
问题——闪烁的混乱
就在几个月前,创建高质量彩色视频的技术似乎遥不可及。如果您使用最初的 DeOldify 模型,并且仅仅像处理任何其他图像一样为每一帧着色,结果就是这样——一片闪烁的混乱
通常首先想到的直接解决方案是某种时序建模,即在训练过程中对模型施加帧窗口的约束,以保持上色决策更加一致。这似乎说得通,但这确实增加了大量的复杂性,并且场景变化的罕见情况引发了关于如何处理连续性的进一步问题。问题很快就变得越来越复杂。基于需要强制时序连贯性的这些假设,制作无缝且无闪烁的彩色视频的前景似乎相当遥远。幸运的是,事实证明这些假设是错误的。
问题在 NoGAN 的帮助下迎刃而解
我们在开发 DeOldify 时发现了一个令人惊讶的观察结果:即使在 GAN 训练之后,并且在不同模型之间,上色决策也并非任意的。相反,不同的模型和训练方案最终都会得到几乎相同的解决方案,只是在上色方面存在非常微小的差异。即使对于那些你可能认为在黑白照片的亮度信息中是不可知且不受约束的事物(例如服装、汽车、电影特效等),情况也是如此。我们还不确定模型究竟学习了什么,才能或多或少地确定性地上色图像。但归根结底是,当你意识到没有什么需要跟踪时,时序连贯性问题就会迎刃而解——帧中的物体无论如何都会以相同的方式渲染。
此外,事实证明,在 NoGAN 训练中,“拐点”之前发生的学习以一种非常有效的方式训练生成器。这不仅体现在快速实现良好的上色效果,还体现在生成器渲染中不引入伪影、变色和不一致性。换句话说,上面《大都会》渲染的“闪烁的混乱”中的那些伪影和故障是由于过多的 GAN 训练造成的,而我们可以通过 NoGAN 几乎完全解决这个问题!
NoGAN 是实现视频渲染稳定性的最重要因素,但还有一些额外的设计选择也会产生影响。例如,更大的 resnet 主干网络 (resnet101) 在特征检测的准确性和稳健性方面产生了显著差异,因此在物体和场景移动时,帧的渲染一致性也更高。另一个考虑因素是渲染分辨率——在某些情况下提高渲染分辨率会产生积极影响,但远不如人们预期的那样大。我们渲染的大多数视频的分辨率范围在 224px 到 360px 之间,这通常效果很好。
所有这些努力的最终结果是,通过简单地将每一帧单独渲染,就像处理任何其他图像一样,就实现了无闪烁、时序一致且色彩丰富的视频!其中完全没有涉及到时序建模。
DeOldify 墓地:失败的尝试
每进行一次成功的实验设计,至少有十次是失败的。这份清单并非详尽无遗,但它包含了我们认为特别有帮助的经验教训。
Wasserstein GAN (WGAN) 及其变体
在开发 DeOldify 时尝试的最初方法是将架构基于Wasserstein GAN (WGAN)。事实证明,WGAN 及其后续改进的稳定性不足以用于着色的实际应用。训练在一段时间内可以工作,但会像 GANs 所知的那样,不可避免地偏离稳定解决方案。在某种程度上,结果是令人愉快的。真正效果极好(甚至是第一次尝试就成功)的是将 DeOldify 建模基于自注意力生成对抗网络。
各种其他归一化方案
尝试了以下归一化变体。它们的效果都不如在生成器中结合批量归一化和谱归一化,而在判别器中只使用谱归一化。
生成器中仅使用谱归一化。 训练速度更慢,通常也不稳定。生成器输出端使用批量归一化。 这显著减慢了训练速度,而且似乎没有任何实际益处。生成器中采用权重归一化。 同样,训练速度慢,图像看起来也不太好。然而有趣的是,在进行超分辨率的 NoGAN 训练时,权重归一化似乎效果最好。
其他损失函数
传统(非 GAN)损失函数与 GAN 损失/训练之间的相互作用至关重要,但也棘手。一方面,您必须仔细权衡非 GAN 和 GAN 损失,这似乎只能通过实验得出。NoGAN 的好处在于,相对于其他 GAN 训练方案,这个过程的迭代速度非常快——只需几分钟而不是几天就能看到最终结果。
损失交互的另一个棘手之处在于,你不能随意选择任何非 GAN 损失函数与 GAN 损失并存。事实证明,感知损失(又称特征损失)在此方面经验上效果最好,尤其与 L1 和均方误差(MSE)损失相比。
由于 NoGAN 训练的大部分重点在于预训练,因此似乎特别重要的是要确保在切换到 GAN 之前,预训练在渲染质量方面尽可能深入。感知损失正好做到了这一点——它本身就能产生所有尝试过的非 GAN 损失中最具色彩的结果。相比之下,像 MSE 和 L1 损失这样更简单的损失函数往往会产生暗淡的着色,因为它们鼓励网络“求稳”,默认选择灰色和棕色。
还尝试了对感知损失进行添加。最值得注意的是 Gram 风格损失和 Wasserstein 距离。虽然这两者不能排除,将来会重新探讨,但这些损失在与 GAN 训练结合使用时,会促使出现奇怪的橙色和黄色变色。怀疑是这些损失没有得到有效使用。
模型参数数量减少
关于 DeOldify 的一件让人惊讶的事情是它的模型尺寸很大。在最新迭代中,“视频”和“稳定”模型在解码器端的大部分层上设置了 1000 个滤波器的宽度。“艺术”模型则在标准 DynamicUnet 配置的基础上,将滤波器数量乘以了 1.5。同样,判别器也相当庞大,起始宽度为 256,而非更常见的 64 或 128。许多实验试图减少参数数量,但它们通常都遇到了同一个问题:产生的渲染结果的颜色显著变淡。
创建超分辨率显微镜视频
最后,我们将讨论索尔克研究所用于创建高分辨率显微镜视频的方法的一些细节。创建模型以从低分辨率序列生成高分辨率显微镜视频所涉及步骤的高级概述如下:
- 获取用于训练的高分辨率原始材料和需要改进的低分辨率材料。
- 开发一个
“加渣”函数 。 - 创建由图像元组(3张图像为一组)组成的低分辨率训练数据集。
- 通过对每个图像元组应用两次“加渣”函数,创建两个训练集 A 和 B。
- 使用“稳定性”损失同时在两个训练集上训练模型。
- 使用训练好的模型,通过在真实的低分辨率原始材料上运行它来生成高分辨率视频。
获取原始材料
在索尔克,我们很幸运,因为仅使用合成低分辨率数据就取得了不错的结果。这很重要,因为要获得完美对齐的高分辨率和低分辨率图像对既耗时又罕见——对于视频(活细胞)来说,这甚至更难或不可能。在这种情况下,文件以专有的 czi 格式获取。幸运的是,这里有一个基于 python 的工具可以读取这种格式:here。
开发一个“加渣”函数
为了生成合成训练数据,我们需要一个“加渣”(crappifier)函数。这是一个将高分辨率图像转换为低分辨率图像的函数,该低分辨率图像近似于我们的模型训练好后将要处理的真实低分辨率图像。这个“加渣”函数会注入高斯噪声和泊松噪声两种形式的随机性。这两种噪声都存在于显微镜图像中。我们在设计时受到了CSBDeep团队出色工作的启发。
“加渣”函数可以很简单,但它确实会实质性地影响输出的质量和特性。例如,我们发现如果我们的“加渣”函数向训练数据注入了过多的高频噪声,训练后的模型就会倾向于消除像神经元那样细小精密的结构。
生成用于训练的合成低分辨率数据
下一步是将图像序列及其目标捆绑在一起进行训练,如下所示
这张图片展示了一个训练示例,我们使用 5 张连续图像(
对于电影,我们使用了三张图像的捆绑,并预测了中间时间点对应的高分辨率图像。换句话说,我们使用来自时间
我们选择了 3 张图像,因为这样可以方便地使用现有的为 3 个输入通道编写的超分辨率网络架构、数据加载器和损失函数。
创建第二 组低分辨率数据
为了使用稳定性损失,我们实际上必须对原始材料
使用稳定性损失训练模型
除了我们用于超分辨率的正常损失函数外,我们还需要选择一个稳定性损失的衡量标准。这是衡量将模型应用于两个训练集时输出生成相似性的标准,正如我们之前解释的那样,这两个训练集仅在随机应用噪声方面有所不同。
给定我们将用于预测真实高分辨率图像 T 的低分辨率图像序列 X,我们创建 X1 和 X2,它们是通过分别应用随机噪声生成“加渣”函数得到的。
X1 = crappifier(X) and X2 = crappifier(X)
给定我们训练好的模型 M,我们然后预测 Y1 和 Y2 如下
Y1 = M(X1) and Y2 = M(X2)
得到超分辨率 L1 = loss(Y1, T)
和 L2 = loss(Y2,T)
。我们的稳定性损失是预测图像之间的差异。我们使用了 L1 损失,但你也可以使用特征损失或其他方法来衡量差异
LossStable = loss(Y1,Y2)
因此,我们最终的训练损失为: loss = L1 + L2 + LossStable
生成影片
现在我们已经训练好了一个模型,从低分辨率输入生成高分辨率输出只需将模型应用于一个滑动窗口,在此例中,一次处理三张低分辨率输入图像。使用Imageio是一种便捷的方式来输出多图像 tif 文件或 mp4 文件。
示例