torch.utils.dlpack

创建于: 2018年7月11日 | 最后更新于: 2025年6月13日

torch.utils.dlpack.from_dlpack(ext_tensor) → Tensor

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

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

参数

• ext_tensor (具有 __dlpack__ 属性的对象，或DLPack capsule) –

  要转换的张量或DLPack capsule。

  如果 ext_tensor 是一个张量（或ndarray）对象，它必须支持 __dlpack__ 协议（即，有一个 ext_tensor.__dlpack__ 方法）。否则 ext_tensor 可以是一个DLPack capsule，它是一个不透明的 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 capsule”）。

注意

to_dlpack 是一个遗留的DLPack接口。它返回的capsule在Python中除了作为 from_dlpack 的输入之外，不能用于其他任何用途。更符合习惯的DLPack用法是直接在张量对象上调用 from_dlpack —— 当该对象具有 __dlpack__ 方法时，这就能正常工作，而PyTorch和大多数其他库现在确实都有该方法。

警告

对于使用 to_dlpack 生成的每个capsule，请仅调用 from_dlpack 一次。当一个capsule被多次使用时，其行为是未定义的。

参数

tensor – 要导出的张量

DLPack capsule共享张量的内存。