本篇文章译自英文文档 Getting Starting using TVMC Python: a high-level API for TVM 作者是 Jocelyn Shiue。更多 TVM 中文文档可访问→TVM 中文站

本节将介绍针对 TVM 初学者规划的脚本工具。

开始前假如没有下载示例模型,需求先经过终端下载 resnet 模型:

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

用你了解的文本编辑器来编辑 Python 文件。

第 0 步:导入​

from tvm.driver import tvmc

第 1 步:加载模型

将模型导入 TVMC。这一步将机器学习模型从支撑的结构,转换为 TVM 的高档图形表示语言 —— Relay。这是为 TVM 中的一切模型一致起点。目前支撑的结构:Keras、ONNX、TensorFlow、TFLite 和 PyTorch。

model = tvmc.load('my_model.onnx') # 第 1 步:加载

检查 Relay,可运转 model.summary()。

一切结构都支撑用 shape_dict 参数掩盖输入 shape。关于大多数结构,这是可选的;但对 PyTorch 是必需的,由于 TVM 无法主动查找它。

model = tvmc.load('my_model.onnx', shape_dict={'input1' : [1, 2, 3, 4], 'input2' : [1, 2, 3, 4]}) #第一步: 加载 + shape_dict

引荐经过 netron 检查模型 input/shape_dict。打开模型后,单击第一个节点检查输入部分中的 name 和 shape。

第 2 步:编译

模型现在是用 Relay 表示的,下一步是将其编译到要运转的硬件(称为 target)。这个编译进程将模型从 Relay,翻译成方针机器可理解的底层语言。

编译模型需求一个 tvm.target 字符串。检查 文档 了解有关 tvm.targets 及其选项的更多信息。一些例子如下:

cuda (英伟达 GPU)
llvm (CPU)
llvm -mcpu=cascadelake(英特尔 CPU)

package = tvmc.compile(model, target="llvm") # 第 2 步:编译

编译完结后返回一个 package。

第 3 步:运转

编译后的 package 可在方针硬件上运转。设备输入选项有:CPU、Cuda、CL、Metal 和 Vulkan。

result = tvmc.run(package, device="cpu") # 第 3 步:运转

用 print(result) 打印成果。

第 1.5 步:调优(可选并引荐)

经过调优可进一步进步运转速度。此可选进程用机器学习来检查模型(函数)中的每个操作,并找到一种更快的办法来运转它。这一步经过 cost 模型,以及对可能的 schedule 进行基准化来完结。

这儿的 target 与编译进程用到的 target 是相同的。

tvmc.tune(model, target="llvm") # 第 1.5 步:可选 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 可忽略。调优会使最终成果运转更快,但调优进程会消耗几个小时的时刻。

参看下面的“保存调优成果”部分,若要运用成果,务必将调优成果传给编译。

tvmc.compile(model, target="llvm", tuning_records = "records.log") # 第 2 步:编译

保存并在终端中启动进程

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') # 第 1 步:加载
model.save(desired_model_path)

保存 package

模型编译完结(第 2 步)后,可将 package 保存下来。

tvmc.compile(model, target="llvm", package_path="whatever") # 第 2 步:编译
new_package = tvmc.TVMCPackage(package_path="whatever")
result = tvmc.run(new_package, device="cpu") # 第 3 步:运转

运用 Autoscheduler

运用下一代 TVM,运转速度会更快。schedule 的查找空间曾经是手写的,而现在是主动生成的。 (了解更多: 1,2)

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

保存调优成果

把调优成果保存在文件中,便利以后复用。

  • 办法 1:
log_file = "hello.json"
# 运转 tuning
tvmc.tune(model, target="llvm", tuning_records=log_file)
...
# 运转 tuning,然后复用 tuning 的成果
tvmc.tune(model, target="llvm",tuning_records=log_file)
  • 办法 2:
# 运转 tuning
tuning_records = tvmc.tune(model, target="llvm")
...
# 运转 tuning,然后复用 tuning 的成果
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,)

为远程设备编译模型

为不在本地计算机上的硬件进行编译时,TVMC 支撑运用远程进程调用(RPC)。要设置 RPC 服务器,可参考本 文档 中的“在设备上设置 RPC 服务器”部分。

TVMC 脚本包括以下内容,并进行了相应调整:

tvmc.tune(
     model,
     target=target, # 编译 target 为字符串 // 要编译的设备
     target_host=target_host, # 主机处理器
     hostname=host_ip_address, # 远程基准测验时运用的 RPC 跟踪器的 IP 地址
     port=port_number, # 要连接的 RPC 跟踪器的端口。默以为 9090。
     rpc_key=your_key, # 方针设备的 RPC 跟踪器密钥。提供 rpc_tracker 时需求
)

下载 Python 源代码
下载 Jupyter Notebook