通過 DLPack 構建跨框架深度學習編譯器
深度學習框架，如Tensorflow, PyTorch, and ApacheMxNet，快速原型化和部署深度學習模型提供了強大的工具箱。不幸的是，易用性往往以碎片化為代價：孤立地使用每個框架是很容易的?？v向集成使開發簡化為常用案例，但冒險走出困境可能比較棘手。
一個支持不力的方案是在內存中將算子從一個框架直接傳遞到另一個框架，而沒有任何數據重復或復制。支持此類使用案例，將使用戶能夠將管道串聯在一起，在一個框架（或更快）中，某些算子比在另一個框架中得到更好的支持?？蚣苤g的共享數據表示也將彌補這一差距，并允許編譯器堆棧在為算子生成代碼時針對單一格式。
DLPack是拉伸數據結構的中間內存表示標準。以 DLPack 為共同表示形式，可以利用 TVM 編寫的腳本，這些框架傳統上只能依賴于供應商提供的庫。TVM DLPack功能可以在 DLPack 算子上運行，提供DLPack，將 PyTorch 和 MxNet 等框架中的算子數據結構與零數據拷貝進行銜接。
DLPack 提供了一個簡單、便攜的內存數據結構：
/*!

• \brief Plain C Tensor object, does not manage memory.
  /
  typedef struct {
  /!
  • \brief The opaque data pointer points to the allocated data.
  • This will be CUDA device pointer or cl_mem handle in OpenCL.
  • This pointer is always aligns to 256 bytes as in CUDA.
    /
    void data;
    /! \brief The device context of the tensor /
    DLContext ctx;
    /! \brief Number of dimensions /
    int ndim;
    /! \brief The data type of the pointer/
    DLDataType dtype;
    /! \brief The shape of the tensor /
    int64_t shape;
    /!
  • \brief strides of the tensor,
  • can be NULL, indicating tensor is compact.
    /
    int64_t strides;
    /*! \brief The offset in bytes to the beginning pointer to data */
    uint64_t byte_offset;
    } DLTensor;
    例如，在 TVM 中聲明和編譯矩陣乘法算子，并構建一個使用 DLPack 表示的wrapper ，使該算子能夠支持 PyTorch 算子。還與 MxNet 重復此演示。此擴展允許機器學習開發人員，在不犧牲性能的情況下，將研究代碼快速移植到相對不受支持的硬件平臺。
    DLPack 如何提供框架和 TVM 之間共享的中間wrapper 的插圖：
    

Figure 1
First, we compute a reference output in PyTorch:
import torch
x = torch.rand(56,56)
y = torch.rand(56,56)
z = x.mm(y)
然后，使用默認調度表定義并構建 TVM 矩陣乘法算子：
n = tvm.convert(56)
X = tvm.placeholder((n,n), name=‘X’)
Y = tvm.placeholder((n,n), name=‘Y’)

k = tvm.reduce_axis((0, n), name='k')
Z = tvm.compute((n,n), lambda i,j : tvm.sum(X[i,k]*Y[k,j], axis=k))
s = tvm.create_schedule(Z.op)
fmm = tvm.build(s, [X, Y, Z], target_host='llvm', name='fmm')

簡潔性，不涵蓋 TVM 的大集合原型調度，可以優化矩陣乘法。
然后，將 TVM 功能轉換為支持 PyTorch 算子的功能：
from tvm.contrib.dlpack import to_pytorch_func
# fmm is the previously built TVM function (Python function)
# fmm is the wrapped TVM function (Python function)
fmm_pytorch = to_pytorch_func(fmm)
z2 = torch.empty(56,56)
fmm_pytorch(x, y, z2)
np.testing.assert_allclose(z.numpy(), z2.numpy())
驗證結果是否匹配。
可以重復相同的示例，但使用 MxNet 代替：
import mxnet
from tvm.contrib.mxnet import to_mxnet_func
ctx = mxnet.cpu(0)
x = mxnet.nd.uniform(shape=(56,56), ctx=ctx)
y = mxnet.nd.uniform(shape=(56,56), ctx=ctx)
z = mxnet.nd.empty(shape=(56,56), ctx=ctx)
f = tvm.build(s, [X, Y, Z], target_host=‘llvm’, name=‘f’)
f_mxnet = to_mxnet_func(f)
f_mxnet(x, y, z)
np.testing.assert_allclose(z.asnumpy(), x.asnumpy().dot(y.asnumpy()))
在PyTorch Example示例的hood下
由于TVM提供將dlpack張量轉換為tvm s的功能，反之亦然，因此所有需要的是通過wrapper 功能來增加一些語法。使用帶有dlpack支持的張量框架的通用轉換器，可用于實現簡易轉換器。
NDArrayconvert_functo_pytorch_func
def convert_func(tvm_func, tensor_type, to_dlpack_func):
assert callable(tvm_func)

def _wrapper(*args):args = tuple(ndarray.from_dlpack(to_dlpack_func(arg))\if isinstance(arg, tensor_type) else arg for arg in args)return tvm_func(*args)return _wrapper

def to_pytorch_func(tvm_func):
import torch
import torch.utils.dlpack
return convert_func(tvm_func, torch.Tensor, torch.utils.dlpack.to_dlpack)

總結

以上是生活随笔為你收集整理的通过 DLPack 构建跨框架深度学习编译器的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。