快捷方式

torch.utils.dlpack

torch.utils.dlpack.from_dlpack(ext_tensor) Tensor[source]

将来自外部库的张量转换为 torch.Tensor

返回的 PyTorch 张量将与输入张量共享内存(可能来自另一个库)。请注意,因此就地操作也会影响输入张量的数据。这可能会导致意外问题(例如,其他库可能具有只读标志或不可变数据结构),因此用户应仅在他们确信这是可以的情况下执行此操作。

参数

ext_tensor (具有 __dlpack__ 属性的对象或 DLPack 胶囊) –

要转换的张量或 DLPack 胶囊。

如果 ext_tensor 是一个张量(或 ndarray)对象,它必须支持 __dlpack__ 协议(即,具有 ext_tensor.__dlpack__ 方法)。否则 ext_tensor 可以是 DLPack 胶囊,它是一个不透明的 PyCapsule 实例,通常由 to_dlpack 函数或方法生成。

返回类型

张量

示例

>>> 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 接口。它返回的胶囊除了用作 from_dlpack 的输入外,不能在 Python 中用于任何其他目的。DLPack 的更惯用的用法是直接在张量对象上调用 from_dlpack - 当该对象具有 __dlpack__ 方法时,这将起作用,PyTorch 和大多数其他库现在确实具有该方法。

警告

对于使用 to_dlpack 生成的每个胶囊,仅调用一次 from_dlpack。胶囊被使用多次时的行为是未定义的。

参数

tensor – 要导出的张量

DLPack 胶囊共享张量的内存。

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源