- nnabla.utils.dlpack.from_dlpack(dlp, arr=None)¶
Decode a DLPack to
# 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
from nnabla import NdArray arr = NdArray() from_dlpack(dlp, arr=arr)
dlp (PyCapsule) – A PyCapsule object of a
"dltensor") which internal memory is borrowed by a tensor of an external package. The ownership of the
DLManagedTensoris moved to an
NdArrayobject, 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
NdArrayobject and decodes the DLPack to it.
NdArrayobject borrowing the DLPack tensor.
- Return type
- nnabla.utils.dlpack.to_dlpack(a, dtype=None, ctx=None)¶
Returns a DLPack which owns an internal array object borrowed by a specified
# 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
NdArrayobject. 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
"dltensor") which internal memory is borrowed by the specified
- Return type