torch.utils.dlpack

torch.utils.dlpack.from_dlpack(ext_tensor) → 张量

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

返回的 PyTorch 张量将与输入张量共享内存（该张量可能来自另一个库）。请注意，就地操作也会影响输入张量的数据。这可能会导致意外问题（例如，其他库可能有只读标志或不可变数据结构），因此用户只有在确定这样做没有问题时才应执行此操作。

Parameters

ext_tensor（具有__dlpack__属性的对象，或DLPack胶囊） –

要转换的张量或DLPack胶囊。

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

Return type

张量

示例：

>>> import torch.utils.dlpack
>>> t = torch.arange(4)

# 直接转换张量（支持在 PyTorch >= 1.10 中）
>>> t2 = torch.from_dlpack(t)
>>> t2[:2] = -1  # 显示内存是共享的
>>> t2
tensor([-1, -1,  2,  3])
>>> t
tensor([-1, -1,  2,  3])

# 旧式的 DLPack 用法，带有中间的胶囊对象
>>> capsule = torch.utils.dlpack.to_dlpack(t)
>>> capsule

>>> t3 = torch.from_dlpack(capsule)
>>> t3
tensor([-1, -1,  2,  3])
>>> t3[0] = -9  # 现在我们在 3 个张量之间共享内存
>>> 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 生成的 capsule 只能调用一次 from_dlpack。 当一个 capsule 被多次消费时的行为是未定义的。

Parameters

tensor – 要导出的张量

DLPack 胶囊共享张量的内存。