用 Python 接口编译和优化模型（AutoTVM）#

原作者: Chris Hoge

在 TVMC 教程中，介绍了如何使用 TVM 的命令行界面 TVMC 来编译、运行和微调预训练的视觉模型 ResNet-50 v2。不过，TVM 不仅仅是命令行工具，它也是优化框架，其 API 可用于许多不同的语言，在处理机器学习模型方面给你带来巨大的灵活性。

在本教程中，将涵盖与 TVMC 相同的内容，但展示如何用 Python API 来完成它。完成本节后，将使用 TVM 的 Python API 来完成以下任务：

编译预训练的 ResNet-50 v2 模型供 TVM 运行时使用。
使用编译后的模型，运行真实图像，并解释输出和评估模型性能。
使用 TVM 在 CPU 上调度该模型。
使用 TVM 收集的调度数据重新编译已优化的模型。
通过优化后的模型运行图像，并比较输出和模型的性能。

本节的目的是让你了解 TVM 的能力以及如何通过 Python API 使用它们。

TVM 是一个深度学习编译器框架，有许多不同的模块可用于处理深度学习模型和算子。在本教程中，我们将研究如何使用 Python API 加载、编译和优化一个模型。

首先要导入一些依赖关系，包括用于加载和转换模型的 mxnet，用于下载测试数据的辅助工具，用于处理图像数据的 Python 图像库，用于图像数据预处理和后处理的 numpy，TVM Relay 框架，以及 TVM Graph Executor。

import warnings
from PIL import Image
import numpy as np
import env # 加载 TVM
from tvm.contrib.download import download_testdata
import tvm
from tvm import relay
from tvm.contrib import graph_executor

warnings.filterwarnings('ignore')

下载和加载前端模型#

在本教程中，使用 ResNet-50 v2。ResNet-50 是卷积神经网络，有 50 层深度，旨在对图像进行分类。该模型已经在超过一百万张图片上进行了预训练，有 1000 种不同的分类。该网络的输入图像大小为 224x224。

备注

如果你有兴趣探索更多关于 ResNet-50 模型的结构，建议下载免费的 ML 模型查看器 Netron。

TVM 提供了辅助库来下载预训练的模型。通过该模块提供模型的 URL、文件名和模型类型，TVM 将下载模型并保存到磁盘。

与其他模型格式一起工作

TVM 支持许多流行的模型格式。清单可以在 TVM 文档的编译深度学习模型部分找到。

备注

可以直接使用如下方式下载预训练的模型（以 ONNX 为例）：

model_url = "".join(
    [
        "https://github.com/onnx/models/raw/",
        "master/vision/classification/resnet/model/",
        "resnet50-v2-7.onnx",
    ]
)

model_path = download_testdata(model_url, "resnet50-v2-7.onnx", module="onnx")

MXNet 可直接载入模型：

from mxnet.gluon.model_zoo import vision

model_name = 'resnet50_v2'
gluon_model = vision.get_model(model_name, pretrained=True)

下载、预处理和加载测试图像#

当涉及到预期的张量形状、格式和数据类型时，每个模型都很特别。出于这个原因，大多数模型需要一些预处理和后处理，以确保输入是有效的，并解释输出。TVMC 对输入和输出数据都采用了 NumPy 的 .npz 格式。

作为本教程的输入，将使用一只猫的图像，但你可以自由地用你选择的任何图像来代替这个图像。

https://s3.amazonaws.com/model-server/inputs/kitten.jpg

下载图像数据，然后将其转换成 numpy 数组，作为模型的输入。

img_url = "https://s3.amazonaws.com/model-server/inputs/kitten.jpg"
img_path = download_testdata(img_url, "imagenet_cat.png", module="data")

# resize 到 224x224
with Image.open(img_path) as im:
    resized_image = im.resize((224, 224))

# 转换为 float32
img_data = np.asarray(resized_image).astype("float32")

# 输入图像是在 HWC 布局，而 MXNet 期望 CHW 输入
img_data = np.transpose(img_data, (2, 0, 1))

# 根据 ImageNet 输入规范进行 Normalize
imagenet_mean = np.array([0.485, 0.456, 0.406]).reshape((3, 1, 1))
imagenet_stddev = np.array([0.229, 0.224, 0.225]).reshape((3, 1, 1))
norm_img_data = (img_data / 255 - imagenet_mean) / imagenet_stddev

# 添加批处理维度，设置数据为 4 维 输入：NCHW
img_data = np.expand_dims(norm_img_data, axis=0)

用 Relay 编译模型#

下一步是编译 ResNet 模型。使用 from_mxnet() 导入器将模型导入到 relay。

不同的模型类型，输入的名称可能不同。你可以使用 Netron 这样的工具来检查输入名称。

input_name = "data"
shape_dict = {input_name: img_data.shape}

mod, params = relay.frontend.from_mxnet(gluon_model, shape_dict)

将模型与标准优化一起构建成 TVM 库。

定义正确的目标

指定正确的目标可以对编译后的模块的性能产生巨大影响，因为它可以利用目标上可用的硬件特性。欲了解更多信息，请参考为 x86 CPU 自动调整卷积网络。建议确定你运行的是哪种 CPU，以及可选的功能，并适当地设置目标。例如，对于某些处理器， target = "llvm -mcpu=skylake"，或者对于具有 AVX-512 向量指令集的处理器， target = "llvm-mcpu=skylake-avx512"。

target = "llvm"

with tvm.transform.PassContext(opt_level=3):
    lib = relay.build(mod, target=target, params=params)

One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.

从该库中创建 TVM graph 运行时模块。

dev = tvm.device(str(target), 0)
module = graph_executor.GraphModule(lib["default"](dev))

在 TVM 运行时上执行#

已经编译了模型，下面可以使用 TVM 运行时来进行预测。要使用 TVM 来运行模型并进行预测，需要两样东西：

编译后的模型，也就是我们刚刚制作的模块 module。
对模型的有效输入，以便进行预测。

# dtype = "float32"
module.set_input(input_name, img_data)
module.run()
output_shape = (1, 1000)
# tvm_output = module.get_output(0,
#                                tvm.nd.empty(output_shape)).numpy()
tvm_output = module.get_output(0).numpy()

收集基本性能数据#

想收集一些与这个未优化的模型相关的基本性能数据，并在以后与调整后的模型进行比较。为了帮助说明 CPU 的噪音，在多个批次的重复中运行计算，然后收集一些关于平均值、中位数和标准差的基础统计数据。

import timeit

timing_number = 10
timing_repeat = 10
unoptimized = (
    np.array(timeit.Timer(lambda: module.run()).repeat(repeat=timing_repeat, number=timing_number))
    * 1000
    / timing_number
)
unoptimized = {
    "mean": np.mean(unoptimized),
    "median": np.median(unoptimized),
    "std": np.std(unoptimized),
}

print(unoptimized)

{'mean': 125.14448569039814, 'median': 109.14178050588816, 'std': 46.9623269438434}

对输出进行后处理#

如前所述，每个模型都有自己提供输出张量的特殊方式。

在案例中，需要运行一些后处理，利用为模型提供的查找表，将 ResNet-50 v2 的输出渲染成更适合人类阅读的形式。

from scipy.special import softmax

from gluoncv.data.imagenet.classification import ImageNet1kAttr

# 获取 ImageNet 标签列表
imagenet_1k_attr = ImageNet1kAttr()
labels = imagenet_1k_attr.classes_long
# 获取输出张量
scores = softmax(tvm_output)
scores = np.squeeze(scores)
ranks = np.argsort(scores)[::-1]
for rank in ranks[0:5]:
    print(f"class='{labels[rank]}' with probability={scores[rank]:f}")

class='tiger cat' with probability=0.526644
class='tabby, tabby cat' with probability=0.403282
class='Egyptian cat' with probability=0.036493
class='tiger, Panthera tigris' with probability=0.004262
class='plastic bag' with probability=0.002360

调优模型#

之前的模型是为了在 TVM 运行时工作而编译的，但不包括任何特定平台的优化。在本节中，将向你展示如何使用 TVM 建立针对你工作平台的优化模型。

在某些情况下，当使用编译的模块运行推理时，可能无法获得预期的性能。在这种情况下，可以利用自动调谐器，为模型找到更好的配置，获得性能的提升。TVM 中的调谐是指对模型进行优化以在给定目标上更快地运行的过程。这与训练或微调不同，因为它不影响模型的准确性，而只影响运行时的性能。作为调优过程的一部分，TVM 将尝试运行许多不同的算子实现变体，以观察哪些算子表现最佳。这些运行的结果被储存在调优记录文件中。

在最简单的形式下，调优需要你提供三样东西：

你打算在上面运行这个模型的设备的目标规格
输出文件的路径，调优记录将被存储在该文件中
要调优的模型的路径

import tvm.auto_scheduler as auto_scheduler
from tvm.autotvm.tuner import XGBTuner
from tvm import autotvm

为运行器设置一些基本参数。运行器采用一组特定参数生成的编译代码，并测量其性能。number 指定我们将测试的不同配置的数量，而 repeat 指定我们将对每个配置进行多少次测量。min_repeat_ms 是一个值，指定需要多长时间运行配置测试。如果重复次数低于这个时间，它将被增加。这个选项对于在 GPU 上进行精确的调优是必要的，而对于 CPU 的调优则不需要。把这个值设置为 0 可以禁用它。timeout 为每个测试的配置运行训练代码的时间设置了上限。

number = 10
repeat = 1
min_repeat_ms = 0  # since we're tuning on a CPU, can be set to 0
timeout = 10  # in seconds

# create a TVM runner
runner = autotvm.LocalRunner(
    number=number,
    repeat=repeat,
    timeout=timeout,
    min_repeat_ms=min_repeat_ms,
    enable_cpu_cache_flush=True,
)

创建简单的结构来保存调谐选项。使用 XGBoost 算法来指导搜索。对于生产作业来说，你会想把试验的数量设置得比这里使用的 10 的值大。对于 CPU，推荐 1500，对于 GPU，推荐 3000-4000。所需的试验次数可能取决于特定的模型和处理器，因此值得花一些时间来评估各种数值的性能，以找到调整时间和模型优化之间的最佳平衡。因为运行调谐是需要时间的，我们将试验次数设置为 10 次，但不建议使用这么小的值。early_stopping 参数是在应用提前停止搜索的条件之前，要运行的最小轨数。measure 选项表示将在哪里建立试验代码，以及将在哪里运行。在这种情况下，使用刚刚创建的 LocalRunner 和 LocalBuilder。tuning_records 选项指定了文件来写入调整数据。

tuning_option = {
    "tuner": "xgb",
    "trials": 10,
    "early_stopping": 100,
    "measure_option": autotvm.measure_option(
        builder=autotvm.LocalBuilder(build_func="default"), runner=runner
    ),
    "tuning_records": "resnet-50-v2-autotuning.json",
}

定义调谐搜索算法

默认情况下，这种搜索是使用 XGBoost 网格算法指导的。根据你的模型的复杂性和可用的时间量，你可能想选择一个不同的算法。

设置调谐参数

在这个例子中，为了节省时间，我们将试验次数和提前停止设置为 10。如果你把这些值设置得更高，你可能会看到更多的性能改进，但这是以花时间调整为代价的。收敛所需的试验次数将取决于模型和目标平台的具体情况。

# begin by extracting the tasks from the onnx model
tasks = autotvm.task.extract_from_program(mod["main"], target=target, params=params)

# Tune the extracted tasks sequentially.
for i, task in enumerate(tasks):
    prefix = "[Task %2d/%2d] " % (i + 1, len(tasks))
    tuner_obj = XGBTuner(task, loss_type="rank")
    tuner_obj.tune(
        n_trial=min(tuning_option["trials"], len(task.config_space)),
        early_stopping=tuning_option["early_stopping"],
        measure_option=tuning_option["measure_option"],
        callbacks=[
            autotvm.callback.progress_bar(tuning_option["trials"], prefix=prefix),
            autotvm.callback.log_to_file(tuning_option["tuning_records"]),
        ],
    )

[Task  1/25]  Current/Best:   60.56/ 184.38 GFLOPS | Progress: (10/10) | 10.92 s Done.
[Task  2/25]  Current/Best:    3.48/ 192.73 GFLOPS | Progress: (10/10) | 5.32 s Done.
[Task  3/25]  Current/Best:  131.97/ 131.97 GFLOPS | Progress: (10/10) | 8.48 s Done.
[Task  4/25]  Current/Best:   89.02/ 126.89 GFLOPS | Progress: (10/10) | 5.37 s Done.
[Task  5/25]  Current/Best:  131.30/ 150.86 GFLOPS | Progress: (10/10) | 4.88 s Done.
[Task  6/25]  Current/Best:  111.19/ 192.77 GFLOPS | Progress: (10/10) | 7.98 s Done.
[Task  7/25]  Current/Best:  157.33/ 157.33 GFLOPS | Progress: (10/10) | 4.39 s Done.
[Task  8/25]  Current/Best:   61.16/ 109.41 GFLOPS | Progress: (10/10) | 12.60 s Done.
[Task  9/25]  Current/Best:  161.47/ 205.40 GFLOPS | Progress: (10/10) | 7.61 s Done.
[Task 10/25]  Current/Best:   38.20/ 194.20 GFLOPS | Progress: (10/10) | 5.97 s Done.
[Task 11/25]  Current/Best:   80.76/ 179.60 GFLOPS | Progress: (10/10) | 6.09 s Done.
[Task 12/25]  Current/Best:   82.61/ 121.33 GFLOPS | Progress: (10/10) | 6.33 s Done.
[Task 13/25]  Current/Best:   87.12/ 180.52 GFLOPS | Progress: (10/10) | 6.17 s Done.
[Task 14/25]  Current/Best:   77.79/ 225.42 GFLOPS | Progress: (10/10) | 11.42 s Done.
[Task 15/25]  Current/Best:   99.85/ 156.87 GFLOPS | Progress: (10/10) | 5.48 s Done.
[Task 16/25]  Current/Best:  105.00/ 140.10 GFLOPS | Progress: (10/10) | 5.02 s Done.
[Task 17/25]  Current/Best:   56.20/ 136.63 GFLOPS | Progress: (10/10) | 5.03 s Done.
[Task 18/25]  Current/Best:   42.54/ 121.78 GFLOPS | Progress: (10/10) | 8.11 s Done.
[Task 19/25]  Current/Best:   17.74/ 145.35 GFLOPS | Progress: (10/10) | 6.90 s Done.
[Task 20/25]  Current/Best:   26.30/  99.58 GFLOPS | Progress: (10/10) | 6.03 s Done.
[Task 22/25]  Current/Best:   27.53/ 101.46 GFLOPS | Progress: (10/10) | 6.49 ss Done.
[Task 23/25]  Current/Best:   57.72/  84.81 GFLOPS | Progress: (10/10) | 6.86 s Done.
[Task 25/25]  Current/Best:    0.00/   0.00 GFLOPS | Progress: (0/10) | 0.00 s s Done.
[Task 25/25]  Current/Best:   18.54/  21.77 GFLOPS | Progress: (10/10) | 11.80 s Done.

用调优数据编译优化后的模型#

作为上述调优过程的输出，我们获得了存储在 resnet-50-v2-autotuning.json 的调优记录。编译器将使用这些结果，在你指定的目标上为模型生成高性能代码。

现在，模型的调优数据已经收集完毕，可以使用优化的算子重新编译模型，以加快计算速度。

with autotvm.apply_history_best(tuning_option["tuning_records"]):
    with tvm.transform.PassContext(opt_level=3, config={}):
        lib = relay.build(mod, target=target, params=params)

dev = tvm.device(str(target), 0)
module = graph_executor.GraphModule(lib["default"](dev))

 Done.

验证优化后的模型是否运行并产生相同的结果：

dtype = "float32"
module.set_input(input_name, img_data)
module.run()
output_shape = (1, 1000)
tvm_output = module.get_output(0, tvm.nd.empty(output_shape)).numpy()

scores = softmax(tvm_output)
scores = np.squeeze(scores)
ranks = np.argsort(scores)[::-1]
for rank in ranks[0:5]:
    print("class='%s' with probability=%f" % (labels[rank], scores[rank]))

class='tiger cat' with probability=0.526639
class='tabby, tabby cat' with probability=0.403286
class='Egyptian cat' with probability=0.036493
class='tiger, Panthera tigris' with probability=0.004262
class='plastic bag' with probability=0.002361

比较已调谐和未调谐的模型#

我们想收集一些与这个优化模型相关的基本性能数据，将其与未优化的模型进行比较。根据你的底层硬件、迭代次数和其他因素，你应该看到优化后的模型与未优化的模型相比有性能的提高。

import timeit

timing_number = 10
timing_repeat = 10
optimized = (
    np.array(timeit.Timer(lambda: module.run()).repeat(repeat=timing_repeat, number=timing_number))
    * 1000
    / timing_number
)
optimized = {"mean": np.mean(optimized), "median": np.median(optimized), "std": np.std(optimized)}


print("optimized: %s" % (optimized))
print("unoptimized: %s" % (unoptimized))

optimized: {'mean': 36.29807254000298, 'median': 35.39204624999002, 'std': 2.746931858081105}
unoptimized: {'mean': 56.3440875500055, 'median': 55.918911850039876, 'std': 1.285981826882138}

小结#

在本教程中，我们举了一个简短的例子，说明如何使用 TVM Python API 来编译、运行和调整一个模型。我们还讨论了对输入和输出进行预处理和后处理的必要性。在调优过程之后，我们演示了如何比较未优化和优化后的模型的性能。

这里我们介绍了使用 ResNet-50 v2 本地的简单例子。然而，TVM 支持更多的功能，包括交叉编译、远程执行和剖析/基准测试。

TVM 文档

用 Python 接口编译和优化模型（AutoTVM）

导航

用 Python 接口编译和优化模型（AutoTVM）#

下载和加载前端模型#

下载、预处理和加载测试图像#

用 Relay 编译模型#

在 TVM 运行时上执行#

收集基本性能数据#

对输出进行后处理#

调优模型#

用调优数据编译优化后的模型#

比较已调谐和未调谐的模型#

小结#