DLPack¶
Via a DLPack capsule, you can borrow a tensor from external software as a nnabla.NdArray
and can share a NdArray
to external software.
- nnabla.utils.dlpack.from_dlpack(dlp, arr=None)¶
Decode a DLPack to
NdArray
.Example:
# Create a tensor of an external tool, and encode as an DLPack. import torch from torch.utils.dlpack import to_dlpack t = torch.ones((5, 5), dtype=torch.float32, device=torch.device('cuda')) dlp = to_dlpack(t) # Borrow the DLPack tensor as nnabla.NdArray from nnabla.utils.dlpack import from_dlpack arr = from_dlpack(dlp)
If you want to move an ownership of DLPack to an exiting
NdArray
;from nnabla import NdArray arr = NdArray() from_dlpack(dlp, arr=arr)
- パラメータ
dlp (PyCapsule) -- A PyCapsule object of a
DLManagedTensor
(as"dltensor"
) which internal memory is borrowed by a tensor of an external package. The ownership of theDLManagedTensor
is moved to anNdArray
object, and the PyCapsule object is marked as"used_dltensor"
to inform that the ownership has been moved.arr (NdArray) -- If specified, a given DLPack is decoded to it. Otherwise, it creates a new
NdArray
object and decodes the DLPack to it.
- 戻り値
an
NdArray
object borrowing the DLPack tensor.- 戻り値の型
- nnabla.utils.dlpack.to_dlpack(a, dtype=None, ctx=None)¶
Returns a DLPack which owns an internal array object borrowed by a specified
NdArray
.Example:
# Create a nnabla.NdArray in CUDA. import nnabla as nn from nnabla.ext_utils import get_extension_context ctx = get_extension_context('cudnn') nn.set_default_context(ctx) a = nn.NdArray.from_numpy_array(np.ones((5, 5), dtype=np.float32)) a.cast(np.float32, ctx) # Expose as a DLPack. from nnabla.utils.dlpack import to_dlpack dlp = to_dlpack(a) # Use the DLPack in PyTorch. import torch from torch.utils.dlpack import from_dlpack t = from_dlpack(dlp) # Changing the values in Torch will also be affected in nnabla # because they share memory. t.add_(1) print(a.data) # All values become 2.
- パラメータ
a (NdArray) -- An
NdArray
object. An internal array which is recently modified or created will be encoded into a DLPack.dtype (numpy.dtype) -- If specified, in-place cast operation may be performed before encoding it to a DLPack.
ctx (Context) -- If specified, in-place device transfer operation may be performed before encoding it into a DLPack.
- 戻り値
A PyCapsule object of a
DLManagedTensor
(as"dltensor"
) which internal memory is borrowed by the specifiedNdArray
.- 戻り値の型
PyCapsule