发布者 Khanh LeViet 和 Luiz Gustavo Martins，开发者倡导者

神经风格迁移 是一种优化技术，用于将两张图像（一张内容图像，例如建筑物，和一张风格图像，例如著名画家的作品）混合在一起，使输出图像看起来像内容图像“以参考图像的风格绘制”。今天，我们很高兴分享一个预先训练的风格迁移 TensorFlow Lite 模型，该模型针对移动设备进行了优化，以及一个使用该模型对任何图像进行风格化的 Android 和 iOS 示例应用程序。

在本文中，我们将带您了解如何优化大型 TensorFlow 模型以进行移动部署，以及如何在移动应用程序中使用 TensorFlow Lite 高效地使用它。我们希望您能够使用我们预先训练的风格迁移模型或利用我们的见解来满足您的用例。

背景

风格迁移示例

风格迁移首次发表在 艺术风格的神经算法 中。然而，原始技术在计算上非常昂贵，即使在高端 GPU 上，对图像进行风格化也可能需要几秒钟。几位作者的后续工作（例如）展示了如何加速风格迁移。

在评估了几种模型架构之后，我们决定从 Magenta 的预先训练的任意风格迁移模型开始，用于我们的示例应用程序。该模型可以将任何内容和风格图像作为输入，然后使用前馈神经网络生成风格化的输出图像。与 Gatys 的论文 中的技术相比，该模型允许更快的风格迁移，但它仍然相当大（44 MB）且缓慢（在 Pixel 4 CPU 上为 2340 毫秒）。因此，我们需要优化模型，使其适合在移动应用程序中使用。本文分享了我们在这方面的经验，以及您可以利用的资源。

优化模型架构

我们的风格迁移模型的结构

Magenta 的任意风格迁移模型由两个子网络组成

• 风格预测网络：将风格图像转换为风格嵌入向量。
  
• 风格转换网络：将风格嵌入向量应用于内容图像以生成风格化的图像。
  

Magenta 的风格预测网络有一个 InceptionV3 骨干网络，所以我们用 MobileNetV2 骨干网络代替了它，该网络针对移动设备进行了优化。风格转换网络由多个卷积层组成。我们应用了来自 MobileNet 的宽度乘数思想，将所有卷积层的输出通道数量缩小了 4 倍。

然后，我们必须决定如何训练我们的模型。我们尝试了多种选择：从头开始训练移动模型或从预先训练的 Magenta 的模型中提炼。我们发现，在从头开始优化其他参数时固定 MobileNetV2 的权重可以得到最佳结果。

我们能够在显着缩小和加速模型的同时，实现类似的风格和内容损失水平。

* 在 Pixel 4 CPU 上使用 TensorFlow Lite（带 2 个线程）进行基准测试，2020 年 4 月。 * 查看此 论文，了解有关此风格迁移模型中使用的损失函数定义的更多详细信息

量化

在确定了模型架构后，我们继续使用 量化（使用 TensorFlow 模型优化工具包）进一步缩小我们的移动模型。这是一种重要的技术，适用于大多数 TensorFlow 模型的移动部署，因为它可以将模型大小缩小高达 4 倍，并以微不足道的质量折衷来加快模型推理速度。

在 TensorFlow 提供的量化选项中，我们决定使用 训练后整数量化，因为它在简单性和模型质量之间取得了良好的平衡。在将 TensorFlow 模型转换为 TensorFlow Lite 时，我们只需要提供一小部分训练数据集。

量化后，我们的模型比原始模型小了一个数量级，速度也更快，同时保持了相同水平的风格和内容损失。

* 在 Pixel 4 CPU 上使用 TensorFlow Lite（带 2 个线程）进行基准测试，2020 年 4 月。

部署到移动设备

我们实现了一个 Android 应用程序，以演示如何使用风格迁移模型。该应用程序接收一张风格图像和一张内容图像，并输出一张混合了输入图像的风格和内容的图像。

我们使用手机摄像头通过 Camera2 API 捕获内容图像，并提供一组著名的绘画作品作为风格图像。如上所述，将风格应用于内容图像有两个步骤。首先，我们使用风格预测网络将风格提取为一组浮点数数组。然后，我们使用风格转换网络将这种风格应用于内容图像。

为了在 CPU 和 GPU 上实现最佳性能，我们创建了两组针对每个芯片进行了优化的 TensorFlow Lite 模型。我们使用 int8 量化 模型进行 CPU 推理，使用 float16 量化 模型进行 GPU 推理。GPU 通常比 CPU 具有更好的性能，但它目前只支持 float 模型，而 float 模型比 int8 量化模型更大。以下是 int8 和 float16 模型的性能表现。

* 在 Pixel 4 上使用 TensorFlow Lite 进行基准测试，2020 年 4 月。

另一种可能的性能提升是，如果您只打算在移动应用程序中支持一组固定的风格图像，那么可以缓存风格预测网络的结果。这将使您的应用程序更小，因为您不需要包含风格预测网络，该网络占总网络大小的 91%。这是将过程分成两个模型而不是一个模型的主要原因。

可以在 GitHub 上找到示例，应用风格的主要类是 StyleTransferModelExecutor。

重要的是，我们不要在 UI 线程上运行风格迁移，因为这在计算上很昂贵。相反，我们使用来自 AndroidX 的 ViewModel 类和一个协程，在专用后台线程上运行它，并轻松更新视图。此外，在使用 GPU 代理 运行模型时，TF Lite 解释器初始化、GPU 代理初始化和推理必须都在同一个线程上运行。

生产环境中的风格迁移

Google 艺术与文化应用程序最近添加了 艺术迁移，它使用 TensorFlow Lite 在设备上运行风格迁移。使用的模型与上面的模型非常相似，但优先考虑质量而不是速度和模型大小。如果您有兴趣查看生产环境中的风格迁移，可以尝试一下。

轮到你了

如果您想将风格迁移添加到您自己的应用程序中，可以从下载 移动示例 开始。float16 (预测网络，转换网络) 和 int8 量化版本 (预测网络，转换网络) 这两个模型版本都可以在 TensorFlow Hub 上获得。我们迫不及待地想看看你能创造出什么！别忘了与我们分享您的作品。

资源

在设备上运行机器学习模型具有以下优点：可以保护用户的隐私，同时还能以低延迟实现功能。

在这篇文章中，我们展示了将 TensorFlow 模型直接转换为 TensorFlow Lite 可能只是第一步。为了获得良好的性能，开发人员应该使用量化来优化他们的模型，并在模型质量、模型大小和推理时间之间找到合适的平衡点。

我们使用以下资源创建了我们的模型。它们也可能适用于您在设备上的机器学习用例。

• Magenta 模型库
  Magenta 是一个由 TensorFlow 支持的开源项目。它使用机器学习来创作音乐和艺术。有许多模型可以转换为 TensorFlow Lite，包括这个风格迁移模型。
  
• TensorFlow 模型优化工具包
  模型优化工具包提供了多种方法来优化您的模型，包括量化和剪枝。
  
• TensorFlow Lite 代理
  TensorFlow Lite 可以利用设备上提供的多种硬件加速器类型，包括 GPU 和 DSP，来加快模型推理速度。