诸如 Tensorflow、PyTorch 和 Apache MxNet 等深度学习框架为深度学习的快速原型设计和模型部署提供了强大的工具箱。不幸的是,它们的易用性通常以碎片化为代价:这仅限于单独使用每个框架。垂直整合使得开发流程适用于常见用例,但打破常规可能会非常棘手。
一种支持不足的情况是在内存中直接将张量从一个框架传递到另一个框架,而不发生任何数据复制或拷贝。支持这样的用例将使用户能够将管道串联在一起,其中某些运算符在一个框架中比另一个框架支持得更好(或更快)。框架之间的共享数据表示也将弥补这一差距,并允许编译器堆栈在为运算符生成代码时以单一格式为目标。
DLPack 是张量数据结构的中间内存表示标准。以 DLPack 作为通用表示,我们可以在框架脚本中利用 TVM,而传统上其仅能依赖设备商提供的库。TVM 的打包函数可以在 DLPack 张量上运行,提供从 PyTorch 和 MxNet 等框架桥接张量数据结构并实现零数据复制的封装器。
DLPack 提供了一个简单、可移植的内存数据结构: