一步实现从TF到TF Lite,谷歌提出定制on-device模型框架

近日,谷歌在 Google I/O 发布了 ML Kit,其核心功能之一是「Learn2Compress」技术支持的自动模型压缩服务。Learn2Compress 可直接将 TensorFlow 模型压缩为 TensorFlow Lite 中的设备内置(on-device)模型,可在移动设备上高效运行,而无需担心内存优化和速度问题。

成功的深度学习模型的训练和运行通常需要大量的计算资源、内存和计算能力,这成为其在移动设备和物联网设备上表现良好的障碍。设备内置的机器学习技术使得在移动设备上运行推断成为可能,具有保护数据隐私和随处访问的优势,而无需考虑连接性。设备内置的机器学习系统(如 MobileNet 和 ProjectionNet)通过优化模型效率来解决移动设备上的资源瓶颈。但是,如果希望为自己的个人移动应用程序训练定制的设备内置模型,该怎么办呢?

近日,谷歌在 Google I/O 发布了 ML Kit,使所有移动开发人员都可以利用机器学习。即将推出的 ML Kit 核心功能之一是由我们的研究团队开发的「Learn2Compress」技术支持的自动模型压缩服务。Learn2Compress 支持 TensorFlow Lite 中的自定义设备内置深度学习模型,可在移动设备上高效运行,而无需担心内存优化和速度问题。用于图像分类的 Learn2Compress 将很快可用,研究者可以通过 ML Kit 获取。Learn2Compress 最初将提供给少数开发人员,并在未来几个月里扩大范围。如果希望使用此功能构建自己的模型,可以点击该链接进行注册:https://docs.google.com/forms/d/e/1FAIpQLSd7Uzx6eepXeF5osByifFsBT_L3BJOymIEjG9uz1wa51Fl9dA/viewform。

运行原理

Learn2Compress 是对 ProjectionNet 等之前论文中介绍的学习框架的概括,结合了几种最先进的压缩神经网络模型的技术。它将用户提供的大型预训练 TensorFlow 模型作为输入,执行训练和优化,然后自动生成规模较小、内存效率更高、功耗更低、推断速度更快且准确率损失最小的即用设备内置模型。

一步实现从TF到TF Lite,谷歌提出定制on-device模型框架

Learn2Compress 用于自动生成设备内置机器学习模型。

为此,Learn2Compress 使用了多种神经网络优化和压缩技术,包括:

  • 修剪(pruning):通过删除对于预测结果影响最小的权重或运算(如得分低的权重)来缩小模型。该方法可以达到很好的效果,特别是对于涉及稀疏输入或输出的设备内置模型,这些模型可以被压缩到一半的大小,同时保留 97% 的原始预测质量。

  • 离散化(quantization):该技术在训练过程中特别有用,可以通过减少模型权重和激活值占用的位数提高推断速度。例如,使用 8 位定点表示法替代浮点数可以加速模型推断、减少能耗,并进一步将模型大小压缩到原来的 1/4。

  • 联合训练(joint training)和精炼(distillation):该方法使用老师-学生的学习策略,即使用较大的老师网络(该案例中是用户提供的 TensorFlow 模型)来训练一个紧凑的学生网络(设备内置模型),确保最小的准确率损失。

一步实现从TF到TF Lite,谷歌提出定制on-device模型框架

使用联合训练和精炼方法学习紧凑的学生网络。

老师网络可以被固定(正如在精炼过程中)或联合优化,甚至同时训练多个不同大小的学生网络。因此,Learn2Compress 可以单次生成多个设备内置模型而不是一个,这些模型的大小和推断速度互不相同,开发者可以在其中选取最适合应用需求的模型。这些方法以及迁移学习等技术让压缩过程更加高效,并可更好地扩展到大规模数据集上。

性能如何?

为了展示 Learn2Compress 的有效性,谷歌研究者使用它构建多个图像和自然语言任务中当前最先进深度神经网络(如 MobileNet、NASNet、Inception、ProjectionNet)的紧凑设备内置模型。对于给定任务和数据集,谷歌生成多个不同推断速度和模型大小的设备内置模型。

一步实现从TF到TF Lite,谷歌提出定制on-device模型框架

不同大小的 Learn2Compress 模型和全尺寸基线网络在 CIFAR-10(左)和 ImageNet(右)图像分类任务上的准确率。用于生成 CIFAR-10 和 ImageNet 压缩变量的学生网络分别使用 NASNet 和 MobileNet 的变体架构进行建模。

对于图像分类,Learn2Compress 可以生成适用于移动端应用、具备较好预测准确率的小型快速模型。例如,在 ImageNet 任务上,Learn2Compress 模型的大小是 Inception v3 基线模型的 1/22、MobileNet v1 基线模型的 1/4,而准确率仅下降了 4.6-7%。在 CIFAR-10 上,使用共享参数联合训练多个 Learn2Compress 模型花费的时间仅比训练单个较大 Learn2Compress 模型多 10%,而获得的 3 个压缩模型大小是后者的 1/94,速度是后者的 27 倍,开销是后者的 1/36,且预测质量较好(90-95% 的 top-1 准确率)。

一步实现从TF到TF Lite,谷歌提出定制on-device模型框架

基线模型和 Learn2Compress 模型在 CIFAR-10 图像分类任务上的计算成本和平均预测延迟(Pixel phone)。Learn2Compress 优化的模型使用类似 NASNet 的网络架构。

谷歌很兴奋地看到该模型在开发者用例上的优秀性能。例如,Fishbrain(钓鱼爱好者社交平台)使用 Learn2Compress 将现有图像分类云模型(大小 80MB+,top-3 准确率 91.8%)压缩成规模较小的移动端模型,大小仅有 5MB,而准确率与之前类似。在很多使用案例中,压缩模型的准确率甚至稍微优于原来的较大模型,原因在于前者更好的正则化效应(regularization effect)。

谷歌将继续改进 Learn2Compress,并扩展至图像分类以外的其他任务。谷歌很期待通过云端的 ML Kit 压缩服务实现这一目标。谷歌希望 Learn2Compress 能够帮助开发者更简单地自动构建和优化设备端 ML 模型,以便他们可以集中精力构建强大的 app,创造更酷的用户体验,包括计算机视觉、自然语言处理和其他机器学习应用。

相关推荐