torch.utils.dlpack#
创建于: Jul 11, 2018 | 最后更新于: Jun 13, 2025
- torch.utils.dlpack.from_dlpack(ext_tensor) Tensor[源代码]#
将外部库中的张量转换为
torch.Tensor。返回的 PyTorch 张量将与输入张量共享内存(该张量可能来自另一个库)。请注意,就地操作也将影响输入张量的数据。这可能会导致意外问题(例如,其他库可能有只读标志或不可变数据结构),因此用户应仅在确定这没问题时才执行此操作。
- 参数:
ext_tensor (具有
__dlpack__属性的对象,或 DLPack 胶囊) –要转换的张量或 DLPack 胶囊。
如果
ext_tensor是张量(或 ndarray)对象,则它必须支持__dlpack__协议(即,具有ext_tensor.__dlpack__方法)。否则,ext_tensor可能是一个 DLPack 胶囊,它是一个不透明的PyCapsule实例,通常由to_dlpack函数或方法生成。device (torch.device 或 str 或 None) – 一个可选的 PyTorch 设备,用于指定新张量放置的位置。如果为 None(默认值),则新张量将与
ext_tensor位于同一设备上。copy (bool 或 None) – 一个可选的布尔值,指示是否复制
self。如果为 None,PyTorch 只在必要时复制。
- 返回类型:
示例
>>> 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 ...> >>> 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])
- torch.utils.dlpack.to_dlpack(tensor) PyCapsule#
返回一个代表张量的不透明对象(“DLPack 胶囊”)。
注意
to_dlpack是一个旧的 DLPack 接口。它返回的胶囊在 Python 中除了作为from_dlpack的输入外,不能用于任何其他用途。更惯用的 DLPack 用法是直接在张量对象上调用from_dlpack- 当该对象具有__dlpack__方法时,这就可以工作,而 PyTorch 和大多数其他库现在确实都有这个方法。警告
对于使用
to_dlpack生成的每个胶囊,只调用一次from_dlpack。胶囊被多次使用时的行为是未定义的。- 参数:
tensor – 要导出的张量
DLPack 胶囊共享张量的内存。