开始使用 TVMC Python：TVM 的高级 API#

原作者: Jocelyn Shiue

Step 0: 导入#

导入本地 TVM 环境：

import warnings
from tvm.driver import tvmc

warnings.filterwarnings('ignore')

Step 1: 加载模型#

将模型导入到 tvmc 中。这一步将机器学习模型从受支持的框架转换为 TVM 的高级图表示语言 Relay。这将为 TVM 中的所有模型提供一个统一的起点。目前支持的框架有：Keras、ONNX、Tensorflow、TFLite 和 PyTorch。

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

import onnx

model_path = 'resnet50-v2-7.onnx'
onnx_model = onnx.load(model_path)

model = tvmc.load(model_path, shape_dict={"data": [1, 3, 224, 224]}) # Step 1: Load

所有框架都支持使用 shape_dict 参数覆盖输入 shape。对于大多数框架来说，这是可选的，但对于 Pytorch 来说，这是必要的，因为 TVM 不能自动搜索它。

#Step 1: Load + shape_dict
model = tvmc.load(my_model,
                  shape_dict={'input1': [1, 2, 3, 4],
                              'input2': [1, 2, 3, 4]})

小技巧

查看模型的 input/shape_dict 的推荐方法是通过 netron。打开模型后，单击第一个节点，在 inputs 部分查看名称和形状。

如果你想看 Relay，你可以运行：

model.summary() # 输出内容太多，此处已省略

Step 2: 编译#

既然模型已经在 Relay 中，下一步就是将它编译到需要运行的硬件上。这个硬件称为目标（target）。此编译过程将模型从 Relay 转换为目标机器可以理解的较低级语言。

为了编译模型 tvm.target 字符串是必需的。查看文档，了解更多关于 tvm.target 的信息及其选项。一些例子包括：

  1. cuda (Nvidia GPU)

  2. llvm (CPU)

  3. llvm -mcpu=cascadelake (Intel CPU)

# Step 2: Compile
package = tvmc.compile(model, target="llvm")

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

编译步骤返回 package。

Step 3: 运行#

编译后的包现在可以在硬件目标上运行。设备输入选项有：CPU、Cuda、CL、Metal 和 Vulkan。

使用 CUDA，需要：

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

#Step 3: Run
results = tvmc.run(package, device="cpu")

2022-08-30 09:43:06.688 INFO load_module /tmp/tmpmv4iao74/mod.so

也可以打印结果：

print(results)

Execution time summary:
 mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  
  97.9411      93.8668      127.5822     84.3479      12.3553   
               
Output Names:
 ['output_0']

Tune [可选 && 推荐]#

通过调优可以进一步提高运行速度。这个可选步骤使用机器学习来查看模型（函数）中的每个运算，并试图找到更快的方法来运行它。通过成本模型来做到这一点，并对可能的调度进行基准测试。

此处 target 与编译相同。

# Step 1.5: Optional Tune
# 可以是 "cuda"
tvmc.tune(model, target="llvm")

这将使最终结果更快，但可能需要数小时来调优。

请参阅下面的 保存调优结果。如果希望应用调优结果，请确保将调优结果传递到 compile 中。

# Step 2: Compile
tvmc.compile(model,
             target="cuda",
             tuning_records="records.log")

额外的 TVMC 功能#

保存模型#

为了以后更快，加载模型（Step 1）后保存 Relay 版本。然后，模型将出现在您为稍后转换语法保存它的地方。

model = tvmc.load(model_path, shape_dict={"data": [1, 3, 224, 224]}) #Step 1: Load
desired_model_path = 'new_model.onnx'
model.save(desired_model_path)

保存包#

在模型被编译（Step 2）之后，包也可以被保存。

tvmc.compile(model, target="llvm", package_path="whatever")

new_package = tvmc.TVMCPackage(package_path="whatever")
#Step 3: Run
result = tvmc.run(new_package, device='cpu') 

2022-08-30 09:45:18.629 INFO load_module /tmp/tmp6isdkeif/mod.so

使用 Autoscheduler#

使用下一代 tvm 来启用可能更快的运行速度结果。调度的搜索空间是自动生成的，不像之前需要手写。

参见

  1. 博文：引入 Auto-scheduler TVM

  2. 论文：Ansor : Generating High-Performance Tensor Programs for Deep Learning

tvmc.tune(model,
          target="llvm",
          enable_autoscheduler=True)

保存调优结果#

调优结果可以保存在文件中，以便以后重用。

调优更多复杂模型#

你可能注意到 T 的打印像 .........T.T..T..T..T.T.T.T.T.T. 增加了搜索时间范围：

tvmc.tune(model,
          target='cpu',
          trials=10000,
          timeout=10)

为远程设备编译模型#

当您希望为不在本地机器上的硬件进行编译时，远程过程调用（remote procedural call，简称 RPC）非常有用。tvmc 方法支持这一点。要设置 RPC 服务器，请查看交叉编译和 RPC 文档中的“在设备上设置 RPC 服务器”一节。

在 TVMC 脚本中包括以下内容并进行相应调整：

tvmc.tune(
     model,
     target=target, # Compilation target as string // Device to compile for
     target_host=target_host, # Host processor
     hostname=host_ip_address, # The IP address of an RPC tracker, used when benchmarking remotely.
     port=port_number, # The port of the RPC tracker to connect to. Defaults to 9090.
     rpc_key=your_key, # The RPC tracker key of the target device. Required when rpc_tracker is provided
)