DLPack 協定#

DLPack 協定 是一種穩定的記憶體內資料結構，允許在主要框架之間交換多維陣列或張量。它專為跨硬體支援而設計，意味著它允許在 CPU 以外的裝置（例如 GPU）上交換資料。

DLPack 協定已被 選定為 Python 陣列 API 標準，由 Python 資料 API 標準聯盟 為了在 Python 生態系統中的陣列/張量函式庫之間實現裝置感知資料交換。有關該標準的更多資訊，請參閱協定文件，有關 DLPack 的更多資訊，請參閱 DLPack 的 Python 規格。

PyArrow 中 DLPack 的實作#

DLPack 協定的生產端已針對 pa.Array 實作，可用於在 PyArrow 和其他張量函式庫之間交換資料。支援的資料類型為整數、無號整數和浮點數。該協定不支援遺失資料，這意味著具有遺失值的 PyArrow 陣列無法透過 DLPack 協定傳輸。目前，Arrow 的協定實作僅支援 CPU 裝置上的資料。

協定的資料交換語法包括

  1. from_dlpack(x)：使用實作 __dlpack__ 方法的陣列物件，並在共享記憶體的同時建立新陣列。

  2. __dlpack__(self, stream=None) 和 __dlpack_device__：產生具有 DLPack 結構的 PyCapsule，該結構從 from_dlpack(x) 內部呼叫。

PyArrow 實作協定的第二部分（__dlpack__(self, stream=None) 和 __dlpack_device__），因此可以被實作 from_dlpack 的函式庫使用。

範例#

將 PyArrow CPU 陣列轉換為 NumPy 陣列

>>> import pyarrow as pa
>>> array = pa.array([2, 0, 2, 4])
<pyarrow.lib.Int64Array object at 0x121fd4880>
[
2,
0,
2,
4
]

>>> import numpy as np
>>> np.from_dlpack(array)
array([2, 0, 2, 4])

將 PyArrow CPU 陣列轉換為 PyTorch 張量

>>> import torch
>>> torch.from_dlpack(array)
tensor([2, 0, 2, 4])

將 PyArrow CPU 陣列轉換為 JAX 陣列

>>> import jax
>>> jax.numpy.from_dlpack(array)
Array([2, 0, 2, 4], dtype=int32)
>>> jax.dlpack.from_dlpack(array)
Array([2, 0, 2, 4], dtype=int32)