精度延迟两不误,移动端性能新SOTA,谷歌TF开源轻量级EfficientNet

机器之心报道

机器之心编辑部

我们都知道，谷歌在 2019 年 5 月推出了一种用于图像分类的新模型——EfficientNet，该模型在使用更少参数、消耗更少算力的情况下，达到了目前图像分类 SOTA 的准确度。这令人不由自主地联想到，如果能让 EfficientNet 运行在计算资源受限的移动设备上，是否会为移动或物联网等设备开拓出新的应用场景呢？TensorFlow 最新的博客给出了答案。

今天，谷歌在 GitHub 与 TFHub 上同时发布了 EfficientNet-Lite，该模型运行在 TensorFlow Lite 上，且专门针对移动设备 CPU、GPU 以及 EdgeTPU 做了优化。EfficientNet-Lite 为边缘设备带来了 EfficientNet 上强大的性能，并且提供五个不同版本，让用户能够根据自己的应用场景灵活地在低延迟与高精度之间选择。

项目地址：https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/lite

即便是 EfficientNet-Lite4 这个计算量最大的版本，在 ImageNet top-1 上达到 80.4% 分类精度的同时，也能够实时（30ms/image）运行在 Pixel 4 的 CPU 上。EfficientNet-Lite 与目前流行的图像分类模型对应量化版本的性能比较如下图所示：

图 1. 在 Pixel 4 CPU 上开启四个线程时各模型经过 INT8 量化后的性能比较。

量化与异构硬件方面的难题

由于边缘设备的一些特性，使得 EfficientNet 在其上运行面临着很多挑战。

首先是量化问题。由于浮点运算在边缘设备上的支持有限，需要对模型进行量化处理。然而，这一过程往往涉及复杂的量化感知训练（伪量化），或降低训练后量化模型的准确率。

幸运的是，借助 TensorFlow Lite 中提供的训练后量化流程来对模型进行量化处理，能够在最大程度上降低对其准确率的影响。

其次是异构计算设备带来的问题。要想在移动设备 GPU 或 EdgeTPU 这一类大范围的硬件加速器上，运行同样的模型是一件非常有挑战的事。由于这类硬件针对其应用场景进行了高度定制，它们仅在有限的运算操作上具有良好性能。谷歌的工程师就发现，一些特定硬件加速器并不能很好地支持 EfficientNet 中的一些运算。

为解决以上问题，谷歌对原本的 EfficientNet 做了如下改进：

• 去除 squeeze-and-excitation 网络，因为其在异构计算设备上支持欠佳
• 使用 RELU6 替代所有的 swish 激活函数，这一操作极大地提升了训练后量化的质量
• 放缩模型尺寸时固定住其 stem 与 head 模块以减少放缩后模型的尺寸与计算量

利用 TensorFlow 模型优化工具包进行训练后量化

得益于 TF 模型优化工具包，研究者利用整型训练后量化（integer-only post-training quantization）即可以轻松地对模型进行量化，同时准确度也不会有太多损失。这样就可以将模型大小减少至原来的 1/4，推理速度提升至原来的 2 倍。

下图是 EfficientNet-Lite0 浮点模型与其量化版本在准确度和延迟两个方面的对比情况：

在 4 线程 Pixel 4 CPU 上进行基准测试。

研究者还分享了一些有关训练后量化的经验。他们表示，在首次尝试利用训练后量化时，准确度出现了大幅度下降：ImageNet 数据集上的 Top-1 准确度由 75% 降至 46%。之后，研究者发现准确度的大幅下降是由量化输出范围过大造成的。在本质上，量化是对浮点值进行仿射变换（affine transformation），从而拟合 8 位整型的（int8）bucket：

量化图示。

在下面的示例中，输出张量范围由-168 到 204：

这表明，由于很难将大范围的浮点张量拟合至 8 位整型 bucket，准确度可能已经损失太多了。

为了解决这一问题，研究者将 swish 激活函数替换为「约束范围的」激活函数（relu6），因为后者将输出范围限制在 [0, 6]。做出此改变后，量化模型在 ImageNet 数据集上的 Top-1 准确度由之前的 46% 回升至 74.4%。

在数据集上尝试使用 EfficientNet-Lite

现在我们可以看看如何使用 EfficientNet-Lite 在你的数据集上的性能表现。官方博客建议使用 TensorFlow Lite Model Maker，这是一个可以在已有 TensorFlow 模型上应用迁移学习的工具。其中，用户可以使用自己的输入数据，并以 TensorFlow Lite 的形式输出模型。

TensorFlow Lite Model Maker 支持很多模型架构，包括 MobileNetV2 和所有变体版本的 EfficientNet-Lite。以下为使用 EfficientNet-Lite0 进行图像分类的代码，只需要五行就够了。

# Load your custom datasetdata = ImageClassifierDataLoader.from_folder(flower_path)train_data, test_data = data.split(0.9)# Customize the pre-trained TensorFlow modelmodel = image_classifier.create(train_data, model_spec=efficienetnet_lite0_spec)# Evaluate the modelloss, accuracy = model.evaluate(test_data)# Export as TensorFlow Lite model.model.export('image_classifier.tflite', 'image_labels.txt')

通过改变 model_spec 参数，你可以很容易地尝试不同的模型。对于小数据集，如 tf_flowers，你可以达到 92% 的准确率，而且只需要几分钟时间，训练 5 个 epoch。如果训练更多 epoch，采用更多数据，或者微调整个模型，效果还能更好。

模型建立好以后，可以将其构建为移动端 app。App 会自动下载在 ImageNet 数据集上预训练过的 EfficientNet-Lite，并存放在 asset 文件件。如果你想定制化自己的模型，可以替换掉该文件夹中的模型。

从截图上可以看到，模型可以实现实时推理速度（>= 30 fps）。

参考链接：https://blog.tensorflow.org/2020/03/higher-accuracy-on-vision-models-with-efficientnet-lite.html?linkId=84432329