今天,谷歌在 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 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')
通过改变 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