用 TVMC 编译和优化模型#

原作者：Leandro Nunes, Matthew Barrett, Chris Hoge

在本节中，将使用 TVMC，即 TVM 命令行驱动程序。TVMC 工具，它暴露了 TVM 的功能，如 auto-tuning、编译、profiling 和通过命令行界面执行模型。

在完成本节内容后，将使用 TVMC 来完成以下任务：

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

本节的目的是让你了解 TVM 和 TVMC 的能力，并为理解 TVM 的工作原理奠定基础。

使用 TVMC#

TVMC 是 Python 应用程序，是 TVM Python 软件包的一部分。当你使用 Python 包安装 TVM 时，你将得到 TVMC 作为命令行应用程序，名为 tvmc。这个命令的位置将取决于你的平台和安装方法。

另外，如果你在 $PYTHONPATH 上将 TVM 作为 Python 模块，你可以通过可执行的 python 模块 python -m tvm.driver.tvmc 访问命令行驱动功能。

为简单起见，本教程将提到 TVMC 命令行使用 tvmc <options>，但同样的结果可以用 python -m tvm.driver.tvmc <options>。

你可以使用帮助页面查看：

!python -m tvm.driver.tvmc --help

usage: tvmc [--config CONFIG] [-v] [--version] [-h]
            {micro,run,tune,compile} ...

TVM compiler driver

options:
  --config CONFIG       configuration json file
  -v, --verbose         increase verbosity
  --version             print the version and exit
  -h, --help            show this help message and exit.

commands:
  {micro,run,tune,compile}
    micro               select micro context.
    run                 run a compiled module
    tune                auto-tune a model
    compile             compile a model.

TVMC - TVM driver command-line interface

tvmc 可用的 TVM 的主要功能来自子命令 compile 和 run，以及 tune。要了解某个子命令下的具体选项，请使用 tvmc <subcommand> --help。将在本教程中逐一介绍这些命令，但首先需要下载预训练模型来使用。

获得模型#

在本教程中，将使用 ResNet-50 v2。ResNet-50 是卷积神经网络，有 50 层深度，设计用于图像分类。将使用的模型已经在超过一百万张图片上进行了预训练，有 1000 种不同的分类。该网络输入图像大小为 224x224。如果你有兴趣探究更多关于 ResNet-50 模型的结构，建议下载 Netron，它免费提供的 ML 模型查看器。

在本教程中，将使用 ONNX 格式的模型。

wget https://github.com/onnx/models/raw/main/vision/classification/resnet/model/resnet50-v2-7.onnx

支持的模型格式

TVMC 支持用 Keras、ONNX、TensorFlow、TFLite 和 Torch 创建的模型。如果你需要明确地提供你所使用的模型格式，请使用选项 tvm.driver.tvmc compile --model-format。

更多信息见 python -m tvm.driver.tvmc compile --help。

为 TVM 添加 ONNX 支持

TVM 依赖于你系统中的 ONNX python 库。你可以使用 pip3 install --user onnx onnxoptimizer 命令来安装 ONNX。如果你有 root 权限并且想全局安装 ONNX，你可以去掉 --user 选项。对 onnxoptimizer 的依赖是可选的，仅用于 onnx>=1.9。

将 ONNX 模型编译到 TVM 运行时中#

一旦下载了 ResNet-50 模型，下一步就是对其进行编译。为了达到这个目的，将使用 tvmc compile。从编译过程中得到的输出是模型的 TAR 包，它被编译成目标平台的动态库。可以使用 TVM 运行时在目标设备上运行该模型。

# 这可能需要几分钟的时间，取决于你的机器
!python -m tvm.driver.tvmc compile \
--target "llvm" \
--input-shapes "data:[1,3,224,224]" \
--output build/resnet50-v2-7-tvm.tar \
params/resnet50-v2-7.onnx

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.

查看 tvmc compile 在 module 中创建的文件：

%%bash
mkdir models
tar -xvf build/resnet50-v2-7-tvm.tar -C models

mod.so
mod.json
mod.params

列出了三个文件：

mod.so 是模型，表示为 C++ 库，可以被 TVM 运行时加载。
mod.json 是 TVM Relay 计算图的文本表示。
mod.params 是包含预训练模型参数的文件。

该 module 可以被你的应用程序直接加载，而 model 可以通过 TVM 运行时 API 运行。

定义正确的 target

指定正确的目标（选项 --target）可以对编译后的模块的性能产生巨大的影响，因为它可以利用目标上可用的硬件特性。

欲了解更多信息，请参考为 x86 CPU 自动调优卷积网络。建议确定你运行的是哪种 CPU，以及可选的功能，并适当地设置目标。

用 TVMC 从编译的模块中运行模型#

已经将模型编译到模块，可以使用 TVM 运行时来进行预测。

TVMC 内置了 TVM 运行时，允许你运行编译的 TVM 模型。为了使用 TVMC 来运行模型并进行预测，需要两样东西：

编译后的模块，我们刚刚生成出来。
对模型的有效输入，以进行预测。

当涉及到预期的张量形状、格式和数据类型时，每个模型都很特别。出于这个原因，大多数模型需要一些预处理和后处理，以确保输入是有效的，并解释输出结果。TVMC 对输入和输出数据都采用了 NumPy 的 .npz 格式。这是得到良好支持的 NumPy 格式，可以将多个数组序列化为文件。

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

输入预处理#

对于 ResNet-50 v2 模型，预期输入是 ImageNet 格式的。下面是为 ResNet-50 v2 预处理图像的脚本例子。

你将需要安装支持的 Python 图像库的版本。你可以使用 pip3 install --user pillow 来满足脚本的这个要求。

#!python ./preprocess.py
from tvm.contrib.download import download_testdata
from PIL import Image
import numpy as np

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

with Image.open(img_path) as im:
    # 缩放图片到 224x224
    resized_image = im.resize((224, 224))
    img_data = np.asarray(resized_image).astype("float32")
# 转换为 ONNX 期望 NCHW 输入
img_data = np.transpose(img_data, (2, 0, 1))
# 归一化到 ImageNet 分布
imagenet_mean = np.array([0.485, 0.456, 0.406])
imagenet_stddev = np.array([0.229, 0.224, 0.225])
norm_img_data = np.zeros(img_data.shape).astype("float32")
for i in range(img_data.shape[0]):
    norm_img_data[i, :, :] = (img_data[i, :, :] / 255 - imagenet_mean[i]) / imagenet_stddev[i]
# 添加 batch 维度
img_data = np.expand_dims(norm_img_data, axis=0)
# 保存预处理后数据（格式为 .npz）
np.savez("build/imagenet_cat", data=img_data)

运行已编译的模块#

有了模型和输入数据，可以运行 TVMC 来做预测：

!python -m tvm.driver.tvmc run \
--inputs build/imagenet_cat.npz \
--output build/predictions.npz \
build/resnet50-v2-7-tvm.tar

2023-03-17 14:53:10.700 INFO load_module /tmp/tmp5xszoh9l/mod.so

回顾一下， .tar 模型文件包括 C++ 库，对 Relay 模型的描述，以及模型的参数。TVMC 包括 TVM 运行时，它可以加载模型并根据输入进行预测。当运行上述命令时，TVMC 会输出新文件，predictions.npz，其中包含 NumPy 格式的模型输出张量。

在这个例子中，在用于编译的同一台机器上运行该模型。在某些情况下，可能想通过 RPC Tracker 远程运行它。要阅读更多关于这些选项的信息，请查看：

!python -m tvm.driver.tvmc run --help

输出后处理#

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

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

下面的脚本显示了后处理的例子，从编译的模块的输出中提取标签。

运行这个脚本应该产生以下输出：

#!python ./postprocess.py
import os.path
import numpy as np
from scipy.special import softmax
from tvm.contrib.download import download_testdata

# 下载标签列表
labels_url = "https://s3.amazonaws.com/onnx-model-zoo/synset.txt"
labels_path = download_testdata(labels_url, "synset.txt", module="data")

with open(labels_path, "r") as f:
    labels = [l.rstrip() for l in f]

output_file = "build/predictions.npz"

# 打开输出并读取输出张量
if os.path.exists(output_file):
    with np.load(output_file) as data:
        scores = softmax(data["output_0"])
        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='n02123045 tabby, tabby cat' with probability=0.610552
class='n02123159 tiger cat' with probability=0.367180
class='n02124075 Egyptian cat' with probability=0.019365
class='n02129604 tiger, Panthera tigris' with probability=0.001273
class='n04040759 radiator' with probability=0.000261

试着用其他图像替换猫的图像，看看 ResNet 模型会做出什么样的预测。

自动调优 ResNet 模型#

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

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

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

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

默认搜索算法需要 xgboost，请参阅下面关于优化搜索算法的详细信息：

pip install xgboost cloudpickle

GPU 版本：

conda install -c conda-forge py-xgboost-gpu
pip install cloudpickle

备注

直接运行调优可能会跑不通：

python -m tvmc tune --target "llvm" \
--output build/resnet50-v2-7-autotuner_records.json \
params/resnet50-v2-7.onnx

参考 issuue 13431 解决 tvmc tune resnet50 ERROR 的问题。

import onnx

onnx_model = onnx.load_model('params/resnet50-v2-7.onnx')
onnx_model.graph.input[0].type.tensor_type.shape.dim[0].dim_value = 1
onnx_model.graph.output[0].type.tensor_type.shape.dim[0].dim_value = 1
onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, 'params/resnet50-v2-7-frozen.onnx')

在这个例子中，如果你为 --target 标志指出更具体的目标，你会看到更好的结果。

TVMC 将对模型的参数空间进行搜索，尝试不同的运算符配置，并选择在你的平台上运行最快的一个。尽管这是基于 CPU 和模型操作的指导性搜索，但仍可能需要几个小时来完成搜索。这个搜索的输出将被保存到 resnet50-v2-7-autotuner_records.json 文件中，以后将被用来编译优化的模型。

定义调优搜索算法

默认情况下，这种搜索是使用 XGBoost Grid 算法引导的。根据你的模型的复杂性和可利用的时间，你可能想选择不同的算法。完整的列表可以通过查阅：

!python -m tvm.driver.tvmc tune --help

对于消费级 Skylake CPU 来说，输出结果将是这样的：

!python -m tvm.driver.tvmc tune \
--target "llvm -mcpu=broadwell" \
--output build/resnet50-v2-7-autotuner_records.json \
params/resnet50-v2-7-frozen.onnx

[Task  1/25]  Current/Best:  272.46/ 493.23 GFLOPS | Progress: (40/40) | 21.48 s Done.
[Task  2/25]  Current/Best:  152.39/ 440.48 GFLOPS | Progress: (40/40) | 15.23 s Done.
[Task  3/25]  Current/Best:  184.53/ 542.38 GFLOPS | Progress: (40/40) | 15.20 s Done.
[Task  4/25]  Current/Best:  241.18/ 407.57 GFLOPS | Progress: (40/40) | 17.54 s Done.
[Task  5/25]  Current/Best:  182.73/ 464.18 GFLOPS | Progress: (40/40) | 15.63 s Done.
[Task  6/25]  Current/Best:  536.16/ 536.16 GFLOPS | Progress: (40/40) | 16.13 s Done.
[Task  7/25]  Current/Best:  214.60/ 392.14 GFLOPS | Progress: (40/40) | 15.98 s Done.
[Task  8/25]  Current/Best:  281.15/ 583.36 GFLOPS | Progress: (40/40) | 19.38 s Done.
[Task  9/25]  Current/Best:  146.96/ 399.98 GFLOPS | Progress: (40/40) | 17.36 s Done.
[Task 10/25]  Current/Best:   60.62/ 403.58 GFLOPS | Progress: (40/40) | 15.26 s Done.
[Task 11/25]  Current/Best:  190.37/ 558.11 GFLOPS | Progress: (40/40) | 15.91 s Done.
[Task 12/25]  Current/Best:  204.62/ 511.79 GFLOPS | Progress: (40/40) | 17.37 s Done.
[Task 13/25]  Current/Best:  199.71/ 448.21 GFLOPS | Progress: (40/40) | 16.20 s Done.
[Task 14/25]  Current/Best:  157.68/ 488.08 GFLOPS | Progress: (40/40) | 17.07 s Done.
[Task 15/25]  Current/Best:  228.61/ 483.70 GFLOPS | Progress: (40/40) | 17.03 s Done.
[Task 16/25]  Current/Best:  149.53/ 461.08 GFLOPS | Progress: (40/40) | 15.08 s Done.
[Task 17/25]  Current/Best:  178.52/ 532.27 GFLOPS | Progress: (40/40) | 15.48 s Done.
[Task 18/25]  Current/Best:   66.78/ 530.63 GFLOPS | Progress: (40/40) | 16.16 s Done.
[Task 19/25]  Current/Best:   44.99/ 436.72 GFLOPS | Progress: (40/40) | 17.55 s Done.
[Task 20/25]  Current/Best:  159.20/ 478.63 GFLOPS | Progress: (40/40) | 18.21 s Done.
[Task 21/25]  Current/Best:  177.36/ 469.23 GFLOPS | Progress: (40/40) | 18.89 s Done.
[Task 22/25]  Current/Best:  384.79/ 439.12 GFLOPS | Progress: (40/40) | 15.84 s Done.
[Task 23/25]  Current/Best:  197.92/ 517.98 GFLOPS | Progress: (40/40) | 17.45 s Done.
[Task 25/25]  Current/Best:    0.80/  52.66 GFLOPS | Progress: (40/40) | 32.85 s Done.
 Done.

调谐会话可能需要很长的时间，所以 tvmc tune 提供了许多选项来定制你的调谐过程，在重复次数方面（例如 --repeat 和 --number），要使用的调优算法等等。

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

作为上述调谐过程的输出，获得了存储在 resnet50-v2-7-autotuner_records.json 的调谐记录。这个文件可以有两种使用方式：

作为进一步调谐的输入（通过 tvmc tune --tuning-records）。
作为对编译器的输入

编译器将使用这些结果来为你指定的目标上的模型生成高性能代码。要做到这一点，可以使用 tvmc compile --tuning-records。

获得更多信息：

!python -m tvm.driver.tvmc compile --help

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

!python -m tvm.driver.tvmc  compile \
--target "llvm" \
--tuning-records build/resnet50-v2-7-autotuner_records.json  \
--output build/resnet50-v2-7-tvm_autotuned.tar \
params/resnet50-v2-7-frozen.onnx

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

!python -m tvm.driver.tvmc run \
--inputs build/imagenet_cat.npz \
--output build/predictions.npz \
build/resnet50-v2-7-tvm_autotuned.tar

2023-03-17 15:50:05.484 INFO load_module /tmp/tmpk_0v6k7d/mod.so

!python postprocess.py

class='n02123045 tabby, tabby cat' with probability=0.610553
class='n02123159 tiger cat' with probability=0.367179
class='n02124075 Egyptian cat' with probability=0.019365
class='n02129604 tiger, Panthera tigris' with probability=0.001273
class='n04040759 radiator' with probability=0.000261

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

TVMC 提供了在模型之间进行基本性能基准测试的工具。你可以指定重复次数，并且 TVMC 报告模型的运行时间（与运行时间的启动无关）。可以粗略了解调谐对模型性能的改善程度。

!python -m tvm.driver.tvmc run \
--inputs build/imagenet_cat.npz \
--output build/predictions.npz  \
--print-time \
--repeat 100 \
build/resnet50-v2-7-tvm_autotuned.tar

2023-03-17 15:52:07.029 INFO load_module /tmp/tmp1nt090vr/mod.so
Execution time summary:
 mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  
  43.1426      43.0816      49.2847      40.7562       1.3496   
               

!python -m tvm.driver.tvmc  run \
--inputs build/imagenet_cat.npz \
--output build/predictions.npz  \
--print-time \
--repeat 100 \
build/resnet50-v2-7-tvm.tar

2023-03-17 15:52:49.358 INFO load_module /tmp/tmpfvn7lje9/mod.so
Execution time summary:
 mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  
  49.7214      48.9426      60.2221      46.6976       2.2708   
               

小结#

在本教程中，介绍了 TVMC，用于 TVM 的命令行驱动。演示了如何编译、运行和调优模型。还讨论了对输入和输出进行预处理和后处理的必要性。在调优过程之后，演示了如何比较未优化和优化后的模型的性能。

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

要想知道还有哪些可用的选项，请看 tvmc --help。

在用 Python 接口编译和优化模型教程中，将使用 Python 接口介绍同样的编译和优化步骤。

用 TVMC 编译和优化模型

目录

用 TVMC 编译和优化模型#

使用 TVMC#

获得模型#

将 ONNX 模型编译到 TVM 运行时中#

用 TVMC 从编译的模块中运行模型#

输入预处理#

运行已编译的模块#

输出后处理#

自动调优 ResNet 模型#

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

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

小结#