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
帮助程序功能来运行演示
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_runtime.GraphModule(lib["default"](tvm.context(target, 0)))
runtime.set_input(input_name, inp)
runtime.run()
return runtime.get_output(0).asnumpy(), runtime
从标签到类名的映射，以验证以下模型的输出是否合理
synset = get_synset()
输出：
File /workspace/.tvm_test_data/data/imagenet1000_clsid_to_human.txt exists, skip.
每个人最喜欢的猫形象进行演示
inp = get_imagenet_input()
输出：
File /workspace/.tvm_test_data/data/cat.png exists, skip.
部署量化的PyTorch模型
首先，演示如何使用PyTorch前端加载由PyTorch量化的深度学习模型。
参考下面的PyTorch静态量化教程，以了解其量化工作流程。 https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html
使用此功能来量化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()
量化，跟踪和运行PyTorch Mobilenet v2模型
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.6/dist-packages/torch/quantization/observer.py:121: 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.6/dist-packages/torch/quantization/observer.py:990: 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方言。
可以打印前端的输出，以查看量化模型的表示方式。
会看到特定于量化的算子，例如qnn.quantize，qnn.dequantize，qnn.requantize和qnn.conv2d等。
input_name = "input" # the input name can be be arbitrary for PyTorch frontend.
input_shapes = [(input_name, (1, 3, 224, 224))]
mod, params = relay.frontend.from_pytorch(script_module, input_shapes)
# print(mod) # comment in to see the QNN IR dump
编译并运行Relay 模块
一旦获得量化的Relay 模块，其余的工作流程与运行浮点模型相同。
在Under the hood下，量化专用算子在编译之前被降低为一系列标准Relay 算子。
tvm_result, rt_mod = run_tvm_model(mod, params, input_name, inp, target="llvm")
比较输出标签
应该看到打印出相同的标签。
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', 'lynx, catamount']
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]))
输出：
132 in 1000 raw floating outputs identical.
衡量性能
给出一个示例，说明如何测量TVM编译模型的性能。
n_repeat = 100 # should be bigger to make the measurement more accurate
ctx = tvm.cpu(0)
ftimer = rt_mod.module.time_evaluator("run", ctx, number=1, repeat=n_repeat)
prof_res = np.array(ftimer().results) * 1e3
print("Elapsed average ms:", np.mean(prof_res))
输出：
Elapsed average ms: 20.436994119999998
笔记
推荐此方法的原因如下：
- 测量是在C ++中完成的，因此没有Python开销
- 包括warm up runs
- 可以使用相同的方法在远程设备（Android等）上进行配置文件。
笔记
除非硬件对快速8位指令有特殊支持，否则量化模型不会比FP32模型快。如果没有快速的8位指令，即使模型本身是8位，TVM也会以16位进行量化卷积。
对于x86，在具有AVX512指令集的CPU上可以实现最佳性能。在这种情况下，TVM将最快可用的8位指令用于给定目标。这包括对VNNI 8位点产品指令（CascadeLake或更高版本）的支持。
此外，以下有关CPU性能的一般技巧同样适用：
- 将环境变量TVM_NUM_THREADS设置为物理核心数number of physical cores
- 选择最适合硬件的目标，例如“ llvm -mcpu = skylake-avx512”或“ llvm -mcpu = cascadelake”（将有更多带有AVX512的CPU）