本文介紹如何將深度學習框架量化的模型加載到 TVM。預量化模型的導入是 TVM 中支持的量化之一。有關 TVM 中量化的更多信息，參閱?此處。
這里演示了如何加載和運行由 PyTorch、MXNet 和 TFLite 量化的模型。加載后，可以在任何 TVM 支持的硬件上運行編譯后的量化模型。
首先，導入必要的包：
from PIL import Image
import numpy as np
import torch
from torchvision.models.quantization import mobilenet as qmobilenetimport tvm
from tvm import relay
from tvm.contrib.download import download_testdata
定義運行 demo 的輔助函數：
def get_transform():import torchvision.transforms as transformsnormalize = 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 calibrationmodel(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.pth0%| | 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 pointReturning 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 秒）
下載 Python 源代碼：deploy_prequantized.py
下載 Jupyter Notebook：deploy_prequantized.ipynb
