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

帮助程序功能来运行演示

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）