快捷方式

NonTensorData

class tensordict.NonTensorData(data: 'Any', _metadata: 'dict | None' = None, _is_non_tensor: 'bool' = True, *, batch_size, device=None, names=None)
property device: device

檢索 tensor class 的裝置型別。

dumps(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) T

將 tensordict 儲存到磁碟。

此函式是 memmap() 的代理。

classmethod fields()

返回描述此資料類的欄位的元組。

接受資料類或其例項。元組元素型別為 Field。

classmethod from_tensordict(tensordict, non_tensordict=None, safe=True)

用於例項化新的 tensor class 物件的 tensor class 包裝器。

引數:
  • tensordict (TensorDict) – 張量型別字典

  • non_tensordict (dict) – 包含非張量和巢狀 tensor class 物件的字典

get(key: NestedKey, *args, **kwargs)

獲取使用輸入鍵儲存的值。

引數:
  • key (字串, 字串元組) – 要查詢的鍵。如果是字串元組,則相當於鏈式呼叫 getattr。

  • default – 如果在 tensorclass 中未找到鍵,則使用的預設值。

返回:

使用輸入鍵儲存的值

classmethod load(prefix: str | pathlib.Path, *args, **kwargs) T

從磁碟載入 tensordict。

此類方法是 load_memmap() 的代理。

load_(prefix: str | pathlib.Path, *args, **kwargs)

在當前 tensordict 中從磁碟載入 tensordict。

此類方法是 load_memmap_() 的代理。

classmethod load_memmap(prefix: str | pathlib.Path, device: Optional[device] = None, non_blocking: bool = False, *, out: Optional[TensorDictBase] = None) T

從磁碟載入記憶體對映的 tensordict。

引數:
  • prefix (字串資料夾路徑) – 儲存已儲存 tensordict 的資料夾路徑。

  • device (torch.device等效型別, 可選) – 如果提供,資料將非同步轉換為該裝置。支援 “meta” 裝置,在這種情況下,資料不會被載入,而是建立一組空的“meta”張量。這對於在不實際開啟任何檔案的情況下了解模型總大小和結構非常有用。

  • non_blocking (布林值, 可選) – 如果為 True,在將張量載入到裝置後不會呼叫 synchronize。預設為 False

  • out (TensorDictBase, 可選) – 可選的 tensordict,資料將寫入其中。

示例

>>> from tensordict import TensorDict
>>> td = TensorDict.fromkeys(["a", "b", "c", ("nested", "e")], 0)
>>> td.memmap("./saved_td")
>>> td_load = TensorDict.load_memmap("./saved_td")
>>> assert (td == td_load).all()

此方法還允許載入巢狀的 tensordict。

示例

>>> nested = TensorDict.load_memmap("./saved_td/nested")
>>> assert nested["e"] == 0

tensordict 也可以載入到“meta”裝置上,或者作為假張量載入。

示例

>>> import tempfile
>>> td = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}})
>>> with tempfile.TemporaryDirectory() as path:
...     td.save(path)
...     td_load = TensorDict.load_memmap(path, device="meta")
...     print("meta:", td_load)
...     from torch._subclasses import FakeTensorMode
...     with FakeTensorMode():
...         td_load = TensorDict.load_memmap(path)
...         print("fake:", td_load)
meta: TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=meta,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=meta,
    is_shared=False)
fake: TensorDict(
    fields={
        a: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)
load_state_dict(state_dict: dict[str, Any], strict=True, assign=False, from_flatten=False)

嘗試將 state_dict 就地載入到目標 tensorclass 上。

maybe_to_stack()

如果 NonTensorData 物件具有非空的 batch-size,則將其轉換為 NonTensorStack 物件。

memmap(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True) T

將所有張量寫入新 tensordict 中對應的記憶體對映張量上。

引數:
  • prefix (字串) – 儲存記憶體對映張量的目錄字首。目錄樹結構將模仿 tensordict 的結構。

  • copy_existing (布林值) – 如果為 False(預設值),當 tensordict 中的條目已是儲存在磁碟上並有關聯檔案的張量,但未按 prefix 儲存到正確位置時,將引發異常。如果為 True,任何現有張量將被複制到新位置。

關鍵字引數:
  • num_threads (整型, 可選) – 用於寫入記憶體對映張量的執行緒數。預設為 0

  • return_early (布林值, 可選) – 如果為 Truenum_threads>0,該方法將返回 tensordict 的 future 物件。

  • share_non_tensor (布林值, 可選) – 如果為 True,非張量資料將在程序間共享,並且在單個節點內的任何 worker 上進行的寫入操作(例如就地更新或設定)將更新所有其他 worker 上的值。如果非張量葉子數量很高(例如,共享大量非張量資料堆疊),這可能導致 OOM 或類似錯誤。預設為 False

  • existsok (布林值, 可選) – 如果為 False,則如果同一路徑中已存在張量,將引發異常。預設為 True

然後 TensorDict 被鎖定,這意味著任何非就地寫入操作(例如,重新命名、設定或刪除條目)將引發異常。一旦 tensordict 被解鎖,記憶體對映屬性將變為 False,因為不再保證跨程序的身份一致性。

返回:

如果 return_early=False,則返回一個張量儲存在磁碟上的新 tensordict;否則返回一個 TensorDictFuture 例項。

注意

以這種方式序列化對於深度巢狀的 tensordict 可能很慢,因此不建議在訓練迴圈中呼叫此方法。

memmap_(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True) T

將所有張量就地寫入對應的記憶體對映張量上。

引數:
  • prefix (字串) – 儲存記憶體對映張量的目錄字首。目錄樹結構將模仿 tensordict 的結構。

  • copy_existing (布林值) – 如果為 False(預設值),當 tensordict 中的條目已是儲存在磁碟上並有關聯檔案的張量,但未按 prefix 儲存到正確位置時,將引發異常。如果為 True,任何現有張量將被複制到新位置。

關鍵字引數:
  • num_threads (整型, 可選) – 用於寫入記憶體對映張量的執行緒數。預設為 0

  • return_early (布林值, 可選) – 如果為 Truenum_threads>0,該方法將返回 tensordict 的 future 物件。可以使用 future.result() 查詢結果 tensordict。

  • share_non_tensor (布林值, 可選) – 如果為 True,非張量資料將在程序間共享,並且在單個節點內的任何 worker 上進行的寫入操作(例如就地更新或設定)將更新所有其他 worker 上的值。如果非張量葉子數量很高(例如,共享大量非張量資料堆疊),這可能導致 OOM 或類似錯誤。預設為 False

  • existsok (布林值, 可選) – 如果為 False,則如果同一路徑中已存在張量,將引發異常。預設為 True

然後 TensorDict 被鎖定,這意味著任何非就地寫入操作(例如,重新命名、設定或刪除條目)將引發異常。一旦 tensordict 被解鎖,記憶體對映屬性將變為 False,因為不再保證跨程序的身份一致性。

返回:

如果 return_early=False,則返回 self;否則返回一個 TensorDictFuture 例項。

注意

以這種方式序列化對於深度巢狀的 tensordict 可能很慢,因此不建議在訓練迴圈中呼叫此方法。

memmap_like(prefix: Optional[str] = None, copy_existing: bool = False, *, existsok: bool = True, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) T

建立一個無內容的記憶體對映 tensordict,其形狀與原始 tensordict 相同。

引數:
  • prefix (字串) – 儲存記憶體對映張量的目錄字首。目錄樹結構將模仿 tensordict 的結構。

  • copy_existing (布林值) – 如果為 False(預設值),當 tensordict 中的條目已是儲存在磁碟上並有關聯檔案的張量,但未按 prefix 儲存到正確位置時,將引發異常。如果為 True,任何現有張量將被複制到新位置。

關鍵字引數:
  • num_threads (整型, 可選) – 用於寫入記憶體對映張量的執行緒數。預設為 0

  • return_early (布林值, 可選) – 如果為 Truenum_threads>0,該方法將返回 tensordict 的 future 物件。

  • share_non_tensor (布林值, 可選) – 如果為 True,非張量資料將在程序間共享,並且在單個節點內的任何 worker 上進行的寫入操作(例如就地更新或設定)將更新所有其他 worker 上的值。如果非張量葉子數量很高(例如,共享大量非張量資料堆疊),這可能導致 OOM 或類似錯誤。預設為 False

  • existsok (布林值, 可選) – 如果為 False,則如果同一路徑中已存在張量,將引發異常。預設為 True

然後 TensorDict 被鎖定,這意味著任何非就地寫入操作(例如,重新命名、設定或刪除條目)將引發異常。一旦 tensordict 被解鎖,記憶體對映屬性將變為 False,因為不再保證跨程序的身份一致性。

返回:

一個新的 TensorDict 例項,如果 return_early=False,則資料儲存為記憶體對映張量;否則為一個 TensorDictFuture 例項。

注意

這是將一組大型緩衝區寫入磁碟的推薦方法,因為 memmap_() 會複製資訊,這對於大型內容來說可能會很慢。

示例

>>> td = TensorDict({
...     "a": torch.zeros((3, 64, 64), dtype=torch.uint8),
...     "b": torch.zeros(1, dtype=torch.int64),
... }, batch_size=[]).expand(1_000_000)  # expand does not allocate new memory
>>> buffer = td.memmap_like("/path/to/dataset")
memmap_refresh_()

如果記憶體對映的 tensordict 具有 saved_path,則重新整理其內容。

如果未關聯路徑,此方法將引發異常。

save(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) T

將 tensordict 儲存到磁碟。

此函式是 memmap() 的代理。

set(key: NestedKey, value: Any, inplace: bool = False, non_blocking: bool = False)

設定一個新的鍵值對。

引數:
  • key (str, str 元組) – 要設定的鍵名。如果為 str 元組,則等同於連續呼叫 getattr 後再最終呼叫 setattr。

  • value (Any) – 要儲存在張量類中的值。

  • inplace (bool, 可選) – 如果為 True,set 將嘗試原地更新值。如果為 False 或鍵不存在,則值將被簡單地寫入其目標位置。

返回:

self

state_dict(destination=None, prefix='', keep_vars=False, flatten=False) dict[str, Any]

返回一個 state_dict 字典,可用於從張量類儲存和載入資料。

tolist()

如果批大小不為空,則將資料轉換為列表。

如果批大小為空,則返回資料。

unbind(dim: int)

返回一個元組,其中包含沿指定維度解除繫結的索引張量類例項。

結果張量類例項將與初始張量類例項共享儲存。

文件

訪問 PyTorch 的完整開發者文件

檢視文件

教程

獲取適合初學者和高階開發者的深度教程

檢視教程

資源

查詢開發資源並解答您的問題

檢視資源