【TVM帮助文档学习】开始使用TVMC Python: TVM的高级API

本文翻译自Getting Starting using TVMC Python: a high-level API for TVM — tvm 0.9.dev0 documentation

首先我们需要准备好示例使用resnet模型。如果你还没有的话，可以按以下步骤下载: 

mkdir myscripts
cd myscripts
wget https://github.com/onnx/models/raw/main/vision/classification/resnet/model/resnet50-v2-7.onnx
mv resnet50-v2-7.onnx my_model.onnx
touch tvmcpythonintro.py

下面让我们开始使用自己喜欢的编辑器编写python代码

Step 0: 导入库

from tvm.driver import tvmc

Step 1: 加载模型

将模型导入tvmc。这一步是将TVM支持的深度学习框架模型转换为TVM的高级图表示语言Relay。这是tvm中的所有模型的统一起点。TVM目前支持的框架有:Keras、ONNX、Tensorflow、TFLite和PyTorch。

model = tvmc.load('my_model.onnx') #Step 1: Load

如果你想了解模型的Relay表示，可以运行model.summary()查看

所有框架都支持使用shape_dict参数覆盖输入形状。对于大多数框架来说这是可选的，但对于pytorch来说这是必须的，因为TVM无法自动搜索pythorch的输入形状。

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

建议使用netron查看模型的input/shape_dict。在打开模型之后，单击第一个节点以查看模型的输入名称和形状。

Step 2: 编译

现在模型已经使用Relay表示了，下一步是将其编译到期望的硬件上运行。我们把这个硬件称为目标。这个编译过程是将模型从Relay转换为目标机器能够理解的低级语言。

编译模型时必须传入tvm.target，它是一个字符串。了解更多关于tvm.target及其选项请查看tvm.target — tvm 0.9.dev0 documentation。例如:

1. cuda (Nvidia GPU)

2. llvm (CPU)

3. llvm -mcpu=cascadelake (Intel CPU)

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

 编译步骤返回一个 package

Step 3: 运行

编译得到的package可以在目标硬件上运行。输入的device选项可以选择CPU, Cuda, CL, Metal, and Vulkan

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

 可以使用print(result)打印run的结果

Step 1.5: 调优 [可选 & 推荐]

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

target参数与compile相同。

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

执行后终端打印：

[Task  1/13]  Current/Best:   82.00/ 106.29 GFLOPS | Progress: (48/769) | 18.56 s
[Task  1/13]  Current/Best:   54.47/ 113.50 GFLOPS | Progress: (240/769) | 85.36 s
.....

可能存在用户警告（ UserWarnings），我们可以忽略该告警。调优将使最终结果更快，但它可能耗费几个小时。

参见下面的“保存调优结果”一节。如果希望应用调优结果，请确保将调优结果传递给compile。

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

保存并在终端中运行

python my_tvmc_script.py

注意：屏幕打印可能刷新得非常快

结果示例 

Time elapsed for training: 18.99 s
Execution time summary:
mean (ms)   max (ms)   min (ms)   std (ms)
  25.24      26.12      24.89       0.38


Output Names:
['output_0']

TVMC的附带功能

保存模型 

为了后面的步骤执行更快，可以在加载模型(步骤1)后保存为Relay版本。然后，模型将出现在你为稍后转换语法保存它的地方。

model = tvmc.load('my_model.onnx') #Step 1: Load
model.save(desired_model_path)

保存package

在模型被编译(步骤2)之后，得到的package也可以保存。 

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

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

使用Autoscheduler

使用下一代tvm来实现可能更快的运行速度结果。调度策略搜索空间是自动生成的，不像以前需要手写。(更多信息:Introducing TVM Auto-scheduler (a.k.a. Ansor)、https://arxiv.org/abs/2006.06762) 

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

保存调优结果 

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

方法1

log_file = "hello.json"

# Run tuning
tvmc.tune(model, target="llvm", tuning_records=log_file)

...

# Later run tuning and reuse tuning results
tvmc.tune(model, target="llvm",tuning_records=log_file)

 方法2

# Run tuning
tuning_records = tvmc.tune(model, target="llvm")

...

# Later run tuning and reuse tuning results
tvmc.tune(model, target="llvm",tuning_records=tuning_records)

调优一个更复杂的模型

如果你注意到T的打印 .........T.T..T..T..T.T.T.T.T.T.增加搜索时间范围: 

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

 为远程设备编译模型

当您希望为不在本地机器上的硬件编译模型时，远程过程调用(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
)