DLPack構建跨框架的深度學習編譯器
Tensorflow，PyTorch和ApacheMxNet等深度學習框架提供了一個功能強大的工具包，可用于快速進行原型設計和部署深度學習模型。易用性通常是以碎片為代價的：孤立地使用每個框架是很容易的。垂直集成已使常見用例的開發流程簡化了，但是冒險走過的路可能很棘手。
一個支持不佳的方案是將張量直接從一個框架傳遞到內存中的另一個框架，而沒有任何數據重復或復制。支持這種用例使用戶能夠將管道串聯在一起，其中某些算子在一個框架中得到比在另一個框架中得到更好的支持（或更快速）?？蚣苤g共享的數據表示形式也將彌合這一差距，并在為算子生成代碼時，允許編譯器堆棧以單一格式為目標。
DLPack是用于張量數據結構的中間內存表示標準。使用DLPack作為通用表示，傳統上只能依賴供應商提供的庫的框架編寫的腳本中利用TVM。TVM打包函數可以在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的說明：
    

圖1
首先，在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大量的調度原語集合。如果希望使自定義GEMM算子在的硬件設備上快速運行，請參考詳細的教程。
然后，將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示例的幕后
由于TVM提供了將dlpack張量轉換為tvm的功能，NDArray反之亦然，因此，通過wrapper功能，所需的只是一些語法 syntactic sugar 。 convert_func是用于使用具有dlpack支持的張量的框架的通用轉換器，可以用于實現方便的轉換器，例如 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)

總結

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

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