当前位置：首页 > news >正文

用 TVMC 编译和优化模型（2）

news 2025/8/24 11:30:33

文章目录

前言
一、使用 TVMC
二、获得模型
三、将 ONNX 模型编译到 TVM 运行时中
四、TVMC 从编译的模块中运行模型
- 4.1、输入预处理
- 4.2 运行已编译的模块
- 4.3 输出后处理

前言

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

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

为 TVM 运行时编译预训练 ResNet-50 v2 模型。
通过编译后的模型运行真实图像，并解释输出和模型的性能。
使用 TVM 在 CPU 上调优模型。
使用 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 driveroptions:--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 moduletune                auto-tune a modelcompile             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

--2022-04-26 13:07:52--  https://github.com/onnx/models/raw/main/vision/classification/resnet/model/resnet50-v2-7.onnx
Resolving github.com (github.com)... 20.205.243.166
Connecting to github.com (github.com)|20.205.243.166|:443... connected.
HTTP request sent, awaiting response... 302 Found
Location: https://media.githubusercontent.com/media/onnx/models/main/vision/classification/resnet/model/resnet50-v2-7.onnx [following]
--2022-04-26 13:07:53--  https://media.githubusercontent.com/media/onnx/models/main/vision/classification/resnet/model/resnet50-v2-7.onnx
Resolving media.githubusercontent.com (media.githubusercontent.com)... 185.199.111.133, 185.199.108.133, 185.199.110.133, ...
Connecting to media.githubusercontent.com (media.githubusercontent.com)|185.199.111.133|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 102442450 (98M) [application/octet-stream]
Saving to: ‘resnet50-v2-7.onnx’resnet50-v2-7.onnx  100%[===================>]  97.70M  4.51MB/s    in 25s     2022-04-26 13:08:27 (3.89 MB/s) - ‘resnet50-v2-7.onnx’ saved [102442450/102442450]

支持的模型格式

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

为 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" \--output resnet50-v2-7-tvm.tar \../../_models/resnet50-v2-7.onnx

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 model
tar -xvf resnet50-v2-7-tvm.tar -C model

mod.so
mod.json
mod.params

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

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

四、TVMC 从编译的模块中运行模型

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

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

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

4.1、输入预处理

对于 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 npimg_url = "https://s3.amazonaws.com/model-server/inputs/kitten.jpg"
img_path = download_testdata(img_url, "imagenet_cat.png", module="data")# Resize it to 224x224
resized_image = Image.open(img_path).resize((224, 224))
img_data = np.asarray(resized_image).astype("float32")# ONNX expects NCHW input, so convert the array
img_data = np.transpose(img_data, (2, 0, 1))# Normalize according to 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]# Add batch dimension
img_data = np.expand_dims(norm_img_data, axis=0)# Save to .npz (outputs imagenet_cat.npz)
np.savez("imagenet_cat", data=img_data)

4.2 运行已编译的模块

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

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

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

4.3 输出后处理

如前所述，每个模型都会有自己的特定方式来提供输出张量。需要运行一些后处理，利用为模型提供的查找表，将 ResNet-50 v2 的输出渲染成人类可读的形式。下面的脚本显示了后处理的例子，从编译的模块的输出中提取标签。运行这个脚本应该产生以下输出

#!python ./postprocess.py
import os.path
import numpy as npfrom scipy.special import softmaxfrom tvm.contrib.download import download_testdata# Download a list of labels
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 = "predictions.npz"# Open the output and read the output tensor
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("class='%s' with probability=%f" % (labels[rank], scores[rank]))

class='n02123045 tabby, tabby cat' with probability=0.621104
class='n02123159 tiger cat' with probability=0.356378
class='n02124075 Egyptian cat' with probability=0.019712
class='n02129604 tiger, Panthera tigris' with probability=0.001215
class='n04040759 radiator' with probability=0.000262