在 CUDA 上部署已量化模型#

原作者: Wuwei Lin

本文是使用 TVM 进行自动量化的入门教程。自动量化是 TVM 中的量化方式之一。TVM 中量化的更多细节可以在 Quantization Story 找到。在本教程中，将在 ImageNet 上导入 GluonCV 预训练模型到 Relay，接着量化 Relay 模型，然后执行推理。

加载 TVM 库：

import set_env

import mxnet_cuda

import tvm
from tvm import te
from tvm import relay
from tvm.contrib.download import download_testdata

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

准备数据集#

演示如何准备用于量化的校准数据集。

首先下载 ImageNet 的验证集并对数据集进行预处理。

import mxnet as mx
from mxnet import gluon

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

校准数据集应该是可迭代对象。在 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 模型动物园导入模型。

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

量化模型#

在量化时，需要找到每一层的每个权重和中间 feature map 张量的 scale。

对于权重，scale 是直接根据权重值计算的。支持 power2 和 max 两种模式。两种模式都首先在权重张量内找到最大值。在 power2 模式中，最大值被四舍五入到 2 的幂。如果权重和中间特征映射的比例都是 2 的幂，可以利用 bit shifting 进行乘法。这使得它的计算效率更高。在 max 模式下，以最大值作为 scale。在不 rounding 的情况下，max 模式在某些情况下可能有更好的精度。当 scale 不是二的幂时，将使用不动点（fixed point）乘法。

对于中间 feature map，可以通过数据感知（data-aware）量化来找到 scale。数据感知量化将校准数据集作为输入参数。通过最小化量化前后激活分布之间的 KL 散度来计算 scale。或者，也可以使用预定义的 global 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__":
    import os
    os.environ['MXNET_CUDNN_AUTOTUNE_DEFAULT'] = '0'
    os.environ['MXNET_CUDNN_LIB_CHECKING'] = '0'
    main()

[20:38:55] /work/mxnet/src/io/iter_image_recordio_2.cc:177: ImageRecordIOParser2: /home/pc/.tvm_test_data/val_256_q90.rec, use 4 threads for decoding..
WARNING:autotvm:One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.
[20:39:29] /work/mxnet/src/io/iter_image_recordio_2.cc:177: ImageRecordIOParser2: /home/pc/.tvm_test_data/val_256_q90.rec, use 4 threads for decoding..