超神经HyperAI · 2023年08月15日 · 福建

使用 C++ API 部署 TVM 模块

更多 TVM 中文文档可访问 →Apache TVM 是一个端到端的深度学习编译框架,适用于 CPU、GPU 和各种机器学习加速芯片。 | Apache TVM 中文站

使用 C++ API 部署 TVM 模块

apps/howto_deploy 中给出了部署 TVM 模块的示例,执行下面的命令运行该示例:

cd apps/howto_deploy
./run_example.sh

获取 TVM Runtime 库​

唯一要做的是链接到 target 平台中的 TVM runtime。 TVM 给出了一个最小 runtime,它的开销大约在 300K 到 600K 之间,具体值取决于使用模块的数量。大多数情况下,可用 libtvm_runtime.so 文件去构建。

若构建 libtvm_runtime 有困难,可查看 tvm_runtime_pack.cc(集成了 TVM runtime 的所有示例)。用构建系统来编译这个文件,然后将它包含到项目中。

查看 apps 获取在 iOS、Android 和其他平台上,用 TVM 构建的应用示例。

动态库 vs. 系统模块​

TVM 有两种使用编译库的方法,查看 prepare_test_libs.py 了解如何生成库,查看 cpp_deploy.cc 了解如何使用它们。

  • 把库存储为共享库,并动态加载到项目中。
  • 将编译好的库以系统模块模式绑定到项目中。

动态加载更加灵活,能快速加载新模块。系统模块是一种更 static 的方法,可用在动态库加载不可用的地方。


部署到 Android

为 Android Target 构建模型​

针对 Android target 的 Relay 模型编译遵循和 android_rpc 相同的方法,以下代码会保存 Android target 所需的编译输出:

lib.export_library("deploy_lib.so", ndk.create_shared)
with open("deploy_graph.json", "w") as fo:
    fo.write(graph.json())
with open("deploy_param.params", "wb") as fo:
    fo.write(runtime.save_param_dict(params))

deploy_lib.so、deploy_graph.json、deploy_param.params 将转到 Android target。

适用于 Android Target 的 TVM Runtime​

参考 此处 为 Android target 构建 CPU/OpenCL 版本的 TVM runtime。参考这个 Java 示例来了解 Android Java TVM API,以及如何加载和执行模型。


将 TVM 集成到项目中

TVM runtime 具有轻量级和可移植性的特点,有几种方法可将 TVM 集成到项目中。

下文介绍如何将 TVM 作为 JIT 编译器集成到项目中,从而用它在系统上生成函数的方法

DLPack 支持​

TVM 的生成函数遵循 PackedFunc 约定,它是一个可以接受位置参数(包括标准类型,如浮点、整数、字符串)的函数。PackedFunc 采用 DLPack 约定中的 DLTensor 指针。唯一要做的是创建一个对应的 DLTensor 对象。

集成用户自定义的 C++ 数组​

在 C++ 中唯一要做的就是将你的数组转换为 DLTensor,并将其地址作为 DLTensor* 传递给生成的函数。

集成用户自定义的 Python 数组​

针对 Python 对象 MyArray,需要做:

  • 将 _tvm_tcode 字段添加到返回 tvm.TypeCode.ARRAY_HANDLE 的数组中
  • 在对象中支持
    _tvm_handle 属性(以 Python 整数形式返回 DLTensor 的地址)
  • 用 tvm.register_extension 注册这个类
# 示例代码
import tvm

class MyArray(object):
    _tvm_tcode = tvm.TypeCode.ARRAY_HANDLE

 @property
 def _tvm_handle(self):
        dltensor_addr = self.get_dltensor_addr()
 return dltensor_addr

# 将注册的步骤放在单独的文件 mypkg.tvm.py 中
# 根据需要选择性地导入依赖
tvm.register_extension(MyArray)
推荐阅读
关注数
665
内容数
206
链接人工智能新场景
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息