更多 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)