文 / 软件工程师 Renjie Liu
2019 年 5 月,Google 发布了一系列名为 EfficientNet 的图像分类模型,此类模型在降低一个数量级的参数和算力消耗的情况下,实现了最前沿 (SOTA) 的精度。如果 EfficientNet 可以在边缘设备上运行,则将为计算资源受限的移动和 IoT 设备开拓全新的应用场景。
今天,我们很高兴宣布推出 EfficientNet-Lite (GitHub 地址,TFHub 地址),该模型在 TensorFlow Lite 上运行,并已针对 CPU、GPU 和 EdgeTPU 进行性能优化。EfficientNet-Lite 让边缘设备也能利用 EfficientNet 的强大性能,并提供五个不同版本,让用户能够按照自己的需求从低延迟 / 小模型 (EfficientNet-Lite0) 到高精度 (EfficientNet-Lite4) 之间进行灵活选择。
其中计算量最大的版本,纯整型量化 (Integer-Only Quantized EfficientNet-Lite4) 的 EfficientNet-Lite4,在 ImageNet 上可达到 80.4% 的 Top-1 精度,同时可实时运行在 Pixel 4 CPU (30 毫秒 / 图像)上。下图显示量化 EfficientNet-Lite 模型与目前流行图像分类模型的对应的量化版本的性能对比如下:
挑战:量化与异构硬件
边缘设备的特性也为 EfficientNet 带来了一些挑战。
量化: 受限于许多边缘设备对浮点运算的支持,因此需要对模型进行量化处理。但这一过程往往涉及到复杂的量化感知训练过程,或会降低训练后量化模型的准确性。
值得庆幸的是,我们可以利用工具包内的 TensorFlow Lite 训练后量化流程 来量化模型,并最大限度避免精度损失。
异构硬件 :要想在移动 GPU/EdgeTPU 等多种加速器上运行同样的模型是一项挑战。由于这类硬件的高度定制化,加速器通常只在有限的一些算子中有良好性能表现。我们发现,EfficientNet 中的某些算子并未得到一些加速器的良好支持。
为了解决异构硬件问题,我们对原本的 EfficientNets 进行了如下改进:
- 删除 Squeeze-and-Excitation 网络,因在这类设备上支持欠佳
- 用 RELU6 替代所有 swish 激活函数,从而显著提升训练后量化的质量(稍后说明)
- 放大模型时固定住其 stem 与 head 模块,以减少缩放后模型的大小与计算量
利用 TensorFlow 模型优化工具包进行训练后量化
借助 TensorFlow 模型优化工具包,我们可以通过训练后纯整形量化轻松对模型进行量化,而不会损失太多精度(详情请参阅此处)。这使得模型大小减少为原来的 1/4,推理速度提升 2 倍。
下图显示 EfficientNet-Lite0 浮点模型与量化版本在精度和延迟方面的对比情况:
在四线程 Pixel 4 CPU 上进行基准测试
此外,我们还希望分享一些有关训练后量化的经验。在首次尝试使用训练后量化时,我们发现精度出现大幅度下降:ImageNet 数据集上的 Top-1 精度由 75% 降至 46%。
随后,我们发现精度的大幅下降的原因是量化输出范围过大。量化本质上是将浮点值进行仿射变换到 int8 的值域内:
在我们的示例中,输出张量的范围介于 -168 至 204 之间,具体如下图所示:
这表明我们可能损失太多的精度,因为难以将大范围的浮动张量拟合至 int8 域内。
为解决这一问题,我们将 swish 激活 函数替换为 “范围受限” 的激活函数 (RELU6),因为后者将输出范围限制在 [0, 6]。经过这样的改变,量化模型在 ImageNet 上的 Top-1 模型精度从浮点基线 74.4% 回升至 75.1%。
立即试用 EfficientNet-Lite
利用 EfficientNet-Lite 的强大性能帮助您处理数据。我们建议您使用 TensorFlow Lite Model Maker,它可在已有 TensorFlow 模型上应用迁移学习,并且您可使用自己的输入数据,以 TensorFlow Lite 格式输出生成的模型。
TensorFlow Lite Model 支持多个模型架构,包括 MobileNetV2 和 EfficientNet-Lite 的所有版本。以下示例说明如何仅用 5 行代码构建 EfficientNet-Lite0 图像分类模型:
# Load your custom dataset
data = ImageClassifierDataLoader.from_folder (flower_path)
train_data, test_data = data.split (0.9)
# Customize the pre-trained TensorFlow model
model = image_classifier.create (train_data, model_spec=efficienetnet_lite0_spec)
# Evaluate the model
loss, accuracy = model.evaluate (test_data)
# Export as TensorFlow Lite model.
model.export ('image_classifier.tflite', 'image_labels.txt')
在带有花卉分类的 Colab 上 尝试一下。模型切换很简单,只需更改 model_spec
参数。对于 tf_flowers 这样的小数据集,您只需在数分钟内训练 5 个周期即可达到大约 92% 的精度。如果训练更多周期、使用更多数据,或者微调整个模型,则精度还能提升。
接下来,我们用这个模型来构建一个移动应用。您可以从我们的 图像分类 示例开始,这是使用 EfficientNet-Lite 构建的即时可用的移动应用。该应用会使用 Gradle 任务,将在 ImageNet 数据集上已完成预训练的 EfficientNet-Lite 模型,自动下载到 Assets 目录中。如果您想尝试定制化 Model Maker 创建的模型,可以在 Assets 目录内替换掉该模型。
如屏幕截图所示,EfficientNet-Lite 模型可实时运行推理 (>= 30 fps)。
了解更多信息
构建并运行参考应用(详见说明)。试用 TensorFlow Hub 上的 EfficientNet-Lite,然后使用 TensorFlow Lite Model Maker 针对您的任务定制模型。如需了解 TensorFlow Lite 的更多信息,请访问 TF Lite 网站,试用 TensorFlow 模型优化,然后在 TF Hub 网站上探索更多 TensorFlow Lite 模型。
致谢
Renjie Liu、Xunkai Zhang、Tian Lin、Yuqi Li、Mingxing Tan、Khanh LeViet、Chao Mei、Amy Jang、Luiz GUStavo Martins、Yunlu Li、Suharsh Sivakumar、Raziel Alvarez、Lawrence Chan、Jess Kim、Mike Liang、Shuangfeng Li、Sarah Sirajuddin
参考文献
[1]《EfficientNet:卷积神经网络模型缩放的反思》:[1905.11946] EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks
原文:Higher accuracy on vision models with EfficientNet-Lite / 译:TensorFlow 公众号