使用TVMC Python

玛了个玛卡巴卡

已于 2024-03-13 10:58:28 修改

阅读量889

点赞数 15

分类专栏： TVM 文章标签： TVM python

于 2024-03-12 21:47:14 首次发布

本文链接：https://blog.csdn.net/Albdon/article/details/136664011

版权

TVM 专栏收录该内容

5 篇文章 0 订阅

订阅专栏

将使用resnet50 v2模型：

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

tvmcpythonintro.py：

# 第 0 步：导入
from tvm.driver import tvmc

# 第 1 步：加载
model = tvmc.load('my_model.onnx', shape_dict={'data':[1, 3, 224, 224]}) # 指定shape_dict

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

# 第 3 步：运行
result = tvmc.run(package, device="cpu", benchmark=True, repeat=100)
stat_table = result.format_times()
print(stat_table)

# 第 4 步：调优
tuning_records = tvmc.tune(model, target="llvm", tuning_records="tuning_records.json")

# 第 5 步：使用调优结果编译
package_tuned = tvmc.compile(model, target="llvm", tuning_records = tuning_records, package_path="package_tuned.tar")

# 第 6 步：运行调优后的模型
result_tuned = tvmc.run(package_tuned, device="cpu", benchmark=True, repeat=100)
stat_table_tuned = result_tuned.format_times()
print(stat_table_tuned)

第0步：导入

导入TVMC（TVM的命令行驱动程序）。

第1步：加载

将模型导入TVMC，这一步将机器学习模型从支持的框架（ONNX、PyTorch、Karas、TensorFlow和TFLite）转换为TVM的高级图形表示语言Relay。

所有框架都支持用 shape_dict 参数覆盖输入 shape。对于大多数框架，这是可选的；但对 PyTorch 是必需的，因为 TVM 无法自动搜索它。

可以通过model.summary()查看Relay

第2步：编译

这一步将模型从Relay编译为要运行的硬件（称为target）可理解的底层语言。

可以使用package_path参数保存编译结果到文件，一般是.tar格式的文件。

可以使用TVMCPackage函数加载编译结果：

package = tvmc.TVMCPackage(package_path="package.tar") # 从文件中加载编译结果

第3步：运行

编译后的package可在目标硬件上运行，设备输入的选项有：cpu、cuda、cl、metal和vulkan。

benchmark参数指定了是否对模块的执行进行基准测试，启用基准测试将导致模型执行 $1 + re p e a t * n u mb er$ 次。

repeat参数指定了重复运行多少次，需要benchmark设置为True。

number参数指定了每次重复中要测量的运行次数，需要benchmark设置为True。

print(result)打印模型的输出结果。

第4步：调优

通过调优可以进一步提高运行速度。

target参数和编译过程用到的target是相同的。

enable_autoscheduler参数设置为True时，使用autoscheduler而不是autotvm，调优后的模型运行速度会更快，因为autoscheduler自动生成搜索空间，搜索空间更大。

tuning_records参数指定输出调优结果的文件名，一般是.json格式的文件。

prior_records参数可以在之前的基础上继续进行调优，prior_records参数是字符串类型，可以是文件路径也可以是调优结果字符串。

trials 参数指定了在自动调优过程中，要尝试的调度策略的总数量，默认为10000个。

repeat参数用于指定每个调度策略在评估时应该重复运行多少次，默认为1次。

number参数用于指定在每次repeat中，该调度策略应该连续运行多少次，默认为10次。

第5步：使用调优结果编译

tuning_records参数利用调优结果对模型进行编译，tuning_records参数是字符串类型，可以是文件路径也可以是调优结果字符串

第6步：运行调优后的模型

以第3步相同的形式运行模型，观察调优结果。