評價此頁

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])