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共享張量的記憶體。