本文介绍如何将深度学习框架量化的模型加载到 TVM。预量化模型的导入是 TVM 中支持的量化之一。有关 TVM 中量化的更多信息,参阅 此处。
这里演示了如何加载和运行由 PyTorch、MXNet 和 TFLite 量化的模型。加载后,可以在任何 TVM 支持的硬件上运行编译后的量化模型。
首先,导入必要的包:
定义运行 demo 的辅助函数:
从标签到类名的映射,验证模型的输出是否合理:
用猫的图像进行演示:
部署量化的 PyTorch 模型
首先演示如何用 PyTorch 前端加载由 PyTorch 量化的深度学习模型。
参考 PyTorch 静态量化教程,了解量化的工作流程。
用下面的函数来量化 PyTorch 模型。此函数采用浮点模型,并将其转换为 uint8。这个模型是按通道量化的。
从 torchvision 加载预量化、预训练的 Mobilenet v2 模型
之所以选择 mobilenet v2,是因为该模型接受了量化感知训练,而其他模型则需要完整的训练后校准。
输出结果:
量化、跟踪和运行 PyTorch Mobilenet v2 模型
量化和 jit 的详细信息可参考 PyTorch 网站上的教程。
输出结果:
使用 PyTorch 前端将量化的 Mobilenet v2 转换为 Relay-QNN
PyTorch 前端支持将量化的 PyTorch 模型,转换为具有量化感知算子的等效 Relay 模块。将此表示称为 Relay QNN dialect。
若要查看量化模型是如何表示的,可以从前端打印输出。
可以看到特定于量化的算子,例如 qnn.quantize、qnn.dequantize、qnn.requantize 和 qnn.conv2d 等。
编译并运行 Relay 模块
获得量化的 Relay 模块后,剩下的工作流程与运行浮点模型相同。详细信息请参阅其他教程。
在底层,量化特定的算子在编译之前,会被降级为一系列标准 Relay 算子。
输出结果:
比较输出标签
可看到打印出相同的标签。
输出结果:
但由于数字的差异,通常原始浮点输出不应该是相同的。下面打印 mobilenet v2 的 1000 个输出中,有多少个浮点输出值是相同的。
输出结果:
测试性能
以下举例说明如何测试 TVM 编译模型的性能。
输出结果:
备注
推荐这种方法的原因如下:
- 测试是在 C++ 中完成的,因此没有 Python 开销大
- 包括几个准备工作
- 可用相同的方法在远程设备(Android 等)上进行分析。
备注
如果硬件对 INT8 整数的指令没有特殊支持,量化模型与 FP32 模型速度相近。如果没有 INT8 整数的指令,TVM 会以 16 位进行量化卷积,即使模型本身是 8 位。
对于 x86,在具有 AVX512 指令集的 CPU 上可实现最佳性能。这种情况 TVM 对给定 target 使用最快的可用 8 位指令,包括对 VNNI 8 位点积指令(CascadeLake 或更新版本)的支持。
此外,以下一般技巧对 CPU 性能的提升同样适用:
- 将环境变量 TVM_NUM_THREADS 设置为物理 core 的数量
- 为硬件选择最佳 target,例如 “llvm -mcpu=skylake-avx512” 或 “llvm -mcpu=cascadelake”(未来会有更多支持 AVX512 的 CPU)
部署量化的 MXNet 模型
待更新
部署量化的 TFLite 模型
待更新
脚本总运行时长: (1 分 7.374 秒)