在Cuda上部署量化模型

介紹TVM自動量化。自動量化是TVM中的一種量化方式。將在ImageNet上導入一個GluonCV預先訓練的模型到Relay，量化Relay模型，然后執行推理。
import tvm
from tvm import te
from tvm import relay
import mxnet as mx
from tvm.contrib.download import download_testdata
from mxnet import gluon
import logging
import os

batch_size = 1
model_name = “resnet18_v1”
target = “cuda”
dev = tvm.device(target)

準備數據集

將演示如何準備用于量化的校準數據集。首先下載ImageNet的驗證集，對數據集進行預處理。
calibration_rec = download_testdata(
“http://data.mxnet.io.s3-website-us-west-1.amazonaws.com/data/val_256_q90.rec”,
“val_256_q90.rec”,
)

def get_val_data(num_workers=4):
mean_rgb = [123.68, 116.779, 103.939]
std_rgb = [58.393, 57.12, 57.375]

def batch_fn(batch):return batch.data[0].asnumpy(), batch.label[0].asnumpy()img_size = 299 if model_name == "inceptionv3" else 224
val_data = mx.io.ImageRecordIter(path_imgrec=calibration_rec,preprocess_threads=num_workers,shuffle=False,batch_size=batch_size,resize=256,data_shape=(3, img_size, img_size),mean_r=mean_rgb[0],mean_g=mean_rgb[1],mean_b=mean_rgb[2],std_r=std_rgb[0],std_g=std_rgb[1],std_b=std_rgb[2],
)
return val_data, batch_fn

校準數據集應該是一個iterable對象。在Python中將校準數據集定義為生成器對象。僅使用少量樣本進行校準。
calibration_samples = 10

def calibrate_dataset():
val_data, batch_fn = get_val_data()
val_data.reset()
for i, batch in enumerate(val_data):
if i * batch_size >= calibration_samples:
break
data, _ = batch_fn(batch)
yield {“data”: data}
導入模型

使用Relay MxNet前端，從Gluon模型zoo，導入模型。
def get_model():
gluon_model = gluon.model_zoo.vision.get_model(model_name, pretrained=True)
img_size = 299 if model_name == “inceptionv3” else 224
data_shape = (batch_size, 3, img_size, img_size)
mod, params = relay.frontend.from_mxnet(gluon_model, {“data”: data_shape})
return mod, params
量化模型
在量化中，需要找到每個層的權重和中間特征映射張量的比例。
對于權重，將根據權重值直接計算比例。支持兩種模式：power2和max。這兩種模式都首先在權重張量內找到最大值。在power2模式下，最大值向下舍入為2的冪。如果權重和中間特征映射的比例都是二的冪，可以利用移位進行乘法，計算效率更高。在最大模式下，最大值用作scale。如果沒有舍入，“最大”模式可能具有更好的精度。當scale不是二的冪時，將使用定點乘法。
對于中間特征映射，可以通過數據感知量化，找到scale。數據感知量化，將校準數據集作為輸入參數。通過最小化激活前與量化后分布間的KL散度，計算標度。可以使用預定義的全局scale。可以節省校準時間。但準確度可能會受到影響。
def quantize(mod, params, data_aware):
if data_aware:
with relay.quantize.qconfig(calibrate_mode=“kl_divergence”, weight_scale=“max”):
mod = relay.quantize.quantize(mod, params, dataset=calibrate_dataset())
else:
with relay.quantize.qconfig(calibrate_mode=“global_scale”, global_scale=8.0):
mod = relay.quantize.quantize(mod, params)
return mod
運行推理
創建一個Relay VM，構建和執行模型。
def run_inference(mod):
model = relay.create_executor(“vm”, mod, dev, target).evaluate()
val_data, batch_fn = get_val_data()
for i, batch in enumerate(val_data):
data, label = batch_fn(batch)
prediction = model(data)
if i > 10: # only run inference on a few samples in this tutorial
break

def main():
mod, params = get_model()
mod = quantize(mod, params, data_aware=True)
run_inference(mod)

if name == “main”:
main()
輸出:
/workspace/python/tvm/relay/build_module.py:333: DeprecationWarning: Please use input parameter mod (tvm.IRModule) instead of deprecated parameter mod (tvm.relay.function.Function)
DeprecationWarning,

參考鏈接：
http://tvm.apache.org/docs/how_to/deploy_models/deploy_quantized.html

總結

以上是生活随笔為你收集整理的在Cuda上部署量化模型的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。