def _ortvalues_to_torch_tensor_list(ortvalues, device, c_class=False):
if len(ortvalues) == 0:
return tuple()
if "ort" == device.type:
if not hasattr(C, "to_aten_ort_device_tensor"):
raise AttributeError(
"onnxruntime is missing to_aten_ort_device_tensor needed to support device == 'ort'."
)
return tuple(
C.to_aten_ort_device_tensor(ov if c_class else ov._ortvalue)
for ov in ortvalues)
if not isinstance(ortvalues, list):
raise TypeError("ortvalues must be a list not %r." % type(ortvalues))
if c_class:
res = [_from_dlpack(ov.to_dlpack()) for ov in ortvalues]
for i in range(0, len(res)):
if ortvalues[i].element_type() == TensorProto.BOOL:
res[i] = res[i].to(torch.bool)
else:
res = [_from_dlpack(ov._ortvalue.to_dlpack()) for ov in ortvalues]
for i in range(0, len(res)):
if ortvalues[i]._ortvalue.element_type() == TensorProto.BOOL:
res[i] = res[i].to(torch.bool)
return tuple(res)
def from_dlpack(ext_tensor: Any) -> torch.Tensor:
"""from_dlpack(ext_tensor) -> Tensor
Convers a tensor from a external library into a ``torch.Tensor``
by means of the ``__dlpack__`` protocol.
The tensor will share the memory with the object represented
in the DLPack.
.. warning::
Only call from_dlpack once per capsule. Its behavior when used
on the same capsule multiple times is undefined.
Args:
ext_tensor (object with __dlpack__ attribute or DLPack capsule):
The tensor or DLPack capsule to convert.
"""
if hasattr(ext_tensor, '__dlpack__'):
device = ext_tensor.__dlpack_device__()
# device is either CUDA or ROCm, we need to pass the current
# stream
if device[0] in (DLDeviceType.kDLGPU, DLDeviceType.kDLROCM):
stream = torch.cuda.current_stream('cuda:{}'.format(device[1]))
# cuda_stream is the pointer to the stream and it is a public
# attribute, but it is not documented
dlpack = ext_tensor.__dlpack__(stream=stream.cuda_stream)
else:
dlpack = ext_tensor.__dlpack__()
else:
# Old versions just call the converter
dlpack = ext_tensor
return _from_dlpack(dlpack)
def from_dlpack(ext_tensor: Any) -> torch.Tensor:
"""from_dlpack(ext_tensor) -> Tensor
Converts a tensor from an external library into a ``torch.Tensor``.
The returned PyTorch tensor will share the memory with the input tensor
(which may have come from another library). Note that in-place operations
will therefore also affect the data of the input tensor. This may lead to
unexpected issues (e.g., other libraries may have read-only flags or
immutable data structures), so the user should only do this if they know
for sure that this is fine.
Args:
ext_tensor (object with ``__dlpack__`` attribute, or a DLPack capsule):
The tensor or DLPack capsule to convert.
If ``ext_tensor`` is a tensor (or ndarray) object, it must support
the ``__dlpack__`` protocol (i.e., have a ``ext_tensor.__dlpack__``
method). Otherwise ``ext_tensor`` may be a DLPack capsule, which is
an opaque ``PyCapsule`` instance, typically produced by a
``to_dlpack`` function or method.
Examples::
>>> import torch.utils.dlpack
>>> t = torch.arange(4)
# Convert a tensor directly (supported in PyTorch >= 1.10)
>>> t2 = torch.from_dlpack(t)
>>> t2[:2] = -1 # show that memory is shared
>>> t2
tensor([-1, -1, 2, 3])
>>> t
tensor([-1, -1, 2, 3])
# The old-style DLPack usage, with an intermediate capsule object
>>> capsule = torch.utils.dlpack.to_dlpack(t)
>>> capsule
<capsule object "dltensor" at 0x7f6017d14300>
>>> t3 = torch.from_dlpack(capsule)
>>> t3
tensor([-1, -1, 2, 3])
>>> t3[0] = -9 # now we're sharing memory between 3 tensors
>>> t3
tensor([-9, -1, 2, 3])
>>> t2
tensor([-9, -1, 2, 3])
>>> t
tensor([-9, -1, 2, 3])
"""
if hasattr(ext_tensor, '__dlpack__'):
device = ext_tensor.__dlpack_device__()
# device is either CUDA or ROCm, we need to pass the current
# stream
if device[0] in (DLDeviceType.kDLGPU, DLDeviceType.kDLROCM):
stream = torch.cuda.current_stream('cuda:{}'.format(device[1]))
# cuda_stream is the pointer to the stream and it is a public
# attribute, but it is not documented
# The array API specify that the default legacy stream must be passed
# with a value of 1 for CUDA
# https://data-apis.org/array-api/latest/API_specification/array_object.html?dlpack-self-stream-none#dlpack-self-stream-none # NOQA
is_cuda = device[0] == DLDeviceType.kDLGPU
# Since pytorch is not using PTDS by default, lets directly pass
# the legacy stream
stream_ptr = 1 if is_cuda and stream.cuda_stream == 0 else stream.cuda_stream
dlpack = ext_tensor.__dlpack__(stream=stream_ptr)
else:
dlpack = ext_tensor.__dlpack__()
else:
# Old versions just call the converter
dlpack = ext_tensor
return _from_dlpack(dlpack)
def my_to_tensor(dlpack_structure):
return _from_dlpack(dlpack_structure)
