评价此页

torch.from_dlpack#

torch.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.devicestrNone) – 指定新张量放置位置的可选 PyTorch 设备。如果为 None(默认),新张量将与 ext_tensor 位于同一设备上。

  • copy (boolNone) – 一个可选的布尔值,指示是否复制 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])