在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