DLPack構建跨框架的深度學習編譯器
Tensorflow，PyTorch和ApacheMxNet等深度學習框架提供了一個功能強大的工具包，可用于快速進行原型設計和部署深度學習模型。易用性通常是以碎片為代價的：孤立地使用每個框架是很容易的。垂直集成已使常見用例的開發(fā)流程簡化了，但是冒險走過的路可能很棘手。
一個支持不佳的方案是將張量直接從一個框架傳遞到內(nèi)存中的另一個框架，而沒有任何數(shù)據(jù)重復或復制。支持這種用例使用戶能夠?qū)⒐艿来?lián)在一起，其中某些算子在一個框架中得到比在另一個框架中得到更好的支持（或更快速）。框架之間共享的數(shù)據(jù)表示形式也將彌合這一差距，并在為算子生成代碼時，允許編譯器堆棧以單一格式為目標。
DLPack是用于張量數(shù)據(jù)結(jié)構的中間內(nèi)存表示標準。使用DLPack作為通用表示，傳統(tǒng)上只能依賴供應商提供的庫的框架編寫的腳本中利用TVM。TVM打包函數(shù)可以在DLPack張量上運行，提供包裝程序以橋接帶有零數(shù)據(jù)副本的框架（例如PyTorch和MxNet）中的張量數(shù)據(jù)結(jié)構。
DLPack提供了一種簡單的可移植內(nèi)存數(shù)據(jù)結(jié)構：
/*!

• \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重復此演示。此擴展使機器學習開發(fā)人員可以在不犧牲性能的情況下，將代碼快速移植到相對不受支持的硬件平臺上。
    DLPack如何提供框架和TVM之間共享的中間包wrapper的說明：
    

圖1
首先，在PyTorch中計算參考輸出：
import torch
x = torch.rand(56,56)
y = torch.rand(56,56)
z = x.mm(y)
然后，使用默認調(diào)度定義并構建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')

為簡便起見，沒有涵蓋可用于優(yōu)化矩陣乘法的TVM大量的調(diào)度原語集合。如果希望使自定義GEMM算子在的硬件設備上快速運行，請參考詳細的教程。
然后，將TVM函數(shù)轉(zhuǎn)換為支持PyTorch張量的函數(shù)：
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())
并驗證結(jié)果是否匹配。
可以重復相同的示例，但是使用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示例的幕后
由于TVM提供了將dlpack張量轉(zhuǎn)換為tvm的功能，NDArray反之亦然，因此，通過wrapper功能，所需的只是一些語法 syntactic sugar 。 convert_func是用于使用具有dlpack支持的張量的框架的通用轉(zhuǎn)換器，可以用于實現(xiàn)方便的轉(zhuǎn)換器，例如 to_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)

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。