使用 TVM 部署框架预量化模型
备注
单击 此处 下载完整的示例代码
本文介绍如何将深度学习框架量化的模型加载到 TVM。预量化模型的导入是 TVM 中支持的量化之一。有关 TVM 中量化的更多信息，参阅 此处。
这里演示了如何加载和运行由 PyTorch、MXNet 和 TFLite 量化的模型。加载后，可以在任何 TVM 支持的硬件上运行编译后的量化模型。
首先，导入必要的包 ：
from PIL import Image
import numpy as np
import torch
from torchvision.models.quantization import mobilenet as qmobilenet
import tvm
from tvm import relay
from tvm.contrib.download import download_testdata
定义运行 demo 的辅助函数：
def get_transform():
import torchvision.transforms as transforms
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
return transforms.Compose(
[
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]
)
def get_real_image(im_height, im_width):
img_url = "https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true"
img_path = download_testdata(img_url, "cat.png", module="data")
return Image.open(img_path).resize((im_height, im_width))
def get_imagenet_input():
im = get_real_image(224, 224)
preprocess = get_transform()
pt_tensor = preprocess(im)
return np.expand_dims(pt_tensor.numpy(), 0)
def get_synset():
synset_url = "".join(
[
"https://gist.githubusercontent.com/zhreshold/",
"4d0b62f3d01426887599d4f7ede23ee5/raw/",
"596b27d23537e5a1b5751d2b0481ef172f58b539/",
"imagenet1000_clsid_to_human.txt",
]
)
synset_name = "imagenet1000_clsid_to_human.txt"
synset_path = download_testdata(synset_url, synset_name, module="data")
with open(synset_path) as f:
return eval(f.read())
def run_tvm_model(mod, params, input_name, inp, target="llvm"):
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target=target, params=params)
runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](tvm.device(target, 0)))
runtime.set_input(input_name, inp)
runtime.run()
return runtime.get_output(0).numpy(), runtime
从标签到类名的映射，验证模型的输出是否合理：
synset = get_synset()
用猫的图像进行演示：
inp = get_imagenet_input()
部署量化的 PyTorch 模型
首先演示如何用 PyTorch 前端加载由 PyTorch 量化的深度学习模型。
参考 PyTorch 静态量化教程，了解量化的工作流程。
用下面的函数来量化 PyTorch 模型。此函数采用浮点模型，并将其转换为 uint8。这个模型是按通道量化的。
def quantize_model(model, inp):
model.fuse_model()
model.qconfig = torch.quantization.get_default_qconfig("fbgemm")
torch.quantization.prepare(model, inplace=True)
# Dummy calibration
model(inp)
torch.quantization.convert(model, inplace=True)
从 torchvision 加载预量化、预训练的 Mobilenet v2 模型
之所以选择 mobilenet v2，是因为该模型接受了量化感知训练，而其他模型则需要完整的训练后校准。
qmodel = qmobilenet.mobilenet_v2(pretrained=True).eval()
输出结果：
Downloading: "https://download.pytorch.org/models/mobilenet_v2-b0353104.pth" to /workspace/.cache/torch/hub/checkpoints/mobilenet_v2-b0353104.pth
0%| | 0.00/13.6M [00:00<?, ?B/s]
44%|####4 | 6.03M/13.6M [00:00<00:00, 63.2MB/s]
89%|########8 | 12.1M/13.6M [00:00<00:00, 61.4MB/s]
100%|##########| 13.6M/13.6M [00:00<00:00, 66.0MB/s]
量化、跟踪和运行 PyTorch Mobilenet v2 模型
量化和 jit 的详细信息可参考 PyTorch 网站上的教程。
pt_inp = torch.from_numpy(inp)
quantize_model(qmodel, pt_inp)
script_module = torch.jit.trace(qmodel, pt_inp).eval()
with torch.no_grad():
pt_result = script_module(pt_inp).numpy()
输出结果：
/usr/local/lib/python3.7/dist-packages/torch/ao/quantization/observer.py:179: UserWarning: Please use quant_min and quant_max to specify the range for observers. reduce_range will be deprecated in a future release of PyTorch.
reduce_range will be deprecated in a future release of PyTorch."
/usr/local/lib/python3.7/dist-packages/torch/ao/quantization/observer.py:1126: UserWarning: must run observer before calling calculate_qparams. Returning default scale and zero point
Returning default scale and zero point "
使用 PyTorch 前端将量化的 Mobilenet v2 转换为 Relay-QNN
PyTorch 前端支持将量化的 PyTorch 模型，转换为具有量化感知算子的等效 Relay 模块。将此表示称为 Relay QNN dialect。
若要查看量化模型是如何表示的，可以从前端打印输出。
可以看到特定于量化的算子，例如 qnn.quantize、qnn.dequantize、qnn.requantize 和 qnn.conv2d 等。
input_name = "input" # 对于 PyTorch 前端，输入名称可以是任意的。
input_shapes = [(input_name, (1, 3, 224, 224))]
mod, params = relay.frontend.from_pytorch(script_module, input_shapes)
# print(mod) # 打印查看 QNN IR 转储
编译并运行 Relay 模块
获得量化的 Relay 模块后，剩下的工作流程与运行浮点模型相同。详细信息请参阅其他教程。
在底层，量化特定的算子在编译之前，会被降级为一系列标准 Relay 算子。
target = "llvm"
tvm_result, rt_mod = run_tvm_model(mod, params, input_name, inp, target=target)
输出结果：
/workspace/python/tvm/driver/build_module.py:268: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
"target_host parameter is going to be deprecated. "
比较输出标签
可看到打印出相同的标签。
pt_top3_labels = np.argsort(pt_result[0])[::-1][:3]
tvm_top3_labels = np.argsort(tvm_result[0])[::-1][:3]
print("PyTorch top3 labels:", [synset[label] for label in pt_top3_labels])
print("TVM top3 labels:", [synset[label] for label in tvm_top3_labels])
输出结果：
PyTorch top3 labels: ['tiger cat', 'Egyptian cat', 'tabby, tabby cat']
TVM top3 labels: ['tiger cat', 'Egyptian cat', 'tabby, tabby cat']
但由于数字的差异，通常原始浮点输出不应该是相同的。下面打印 mobilenet v2 的 1000 个输出中，有多少个浮点输出值是相同的。
print("%d in 1000 raw floating outputs identical." % np.sum(tvm_result[0] == pt_result[0]))
输出结果：
154 in 1000 raw floating outputs identical.
测试性能
以下举例说明如何测试 TVM 编译模型的性能。
n_repeat = 100 # 为使测试更准确，应选取更大的数值
dev = tvm.cpu(0)
print(rt_mod.benchmark(dev, number=1, repeat=n_repeat))
输出结果：
Execution time summary:
mean (ms) median (ms) max (ms) min (ms) std (ms)
90.3752 90.2667 94.6845 90.0629 0.6087
备注
推荐这种方法的原因如下：
- 测试是在 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 秒）