PromptData¶

class torchrl.data.PromptData(input_ids: 'torch.Tensor', attention_mask: 'torch.Tensor', prompt_rindex: 'torch.Tensor', labels: 'Optional[torch.Tensor]' = None, logits: 'Optional[torch.Tensor]' = None, loss: 'Optional[torch.Tensor]' = None, *, batch_size, device=None, names=None)[source]¶

property batch_size: Size¶

檢索張量類別的批次大小。

傳回:: 批次大小 (torch.Size)

property device: device¶: 檢索張量類別的裝置類型。

classmethod from_dataset(split, dataset_name=None, max_length=550, root_dir=None, from_disk=False, num_workers: Optional[int] = None)[原始碼]¶

從資料集名稱傳回一個 PromptData。

參數:

split (str) – "train" 或 "valid"，取決於所需的資料分割。
dataset_name (str, optional) – 要處理的資料集名稱。預設為 "CarperAI/openai_summarize_comparisons"。
max_length (int, optional) – 資料集序列的最大長度。預設為 550。
root_dir (path, optional) – 儲存資料集的路徑。預設為 "$HOME/.cache/torchrl/data"。
from_disk (bool, optional) – 如果為 True，將使用 datasets.load_from_disk()。否則，將使用 datasets.load_dataset()。預設為 False。
num_workers (int, optional) – 用於 datasets.dataset.map() 的 workers 數量，該函數在 tokenization 期間被呼叫。預設為 max(os.cpu_count() // 2, 1)。

傳回值: 一個包含記憶體映射版本的必要資料集的 PromptData 實例。: 必要資料集的記憶體映射版本。

範例

>>> data = PromptData.from_dataset("train")
>>> print(data)
PromptDataTLDR(
    attention_mask=MemoryMappedTensor(shape=torch.Size([116722, 550]), device=cpu, dtype=torch.int64, is_shared=False),
    input_ids=MemoryMappedTensor(shape=torch.Size([116722, 550]), device=cpu, dtype=torch.int64, is_shared=False),
    prompt_rindex=MemoryMappedTensor(shape=torch.Size([116722]), device=cpu, dtype=torch.int64, is_shared=False),
    labels=MemoryMappedTensor(shape=torch.Size([116722, 550]), device=cpu, dtype=torch.int64, is_shared=False),
    logits=None,
    loss=None,
    batch_size=torch.Size([116722]),
    device=None,
    is_shared=False)
>>> # data can be sampled from using regular indexing
>>> sub_data = data[:3]

classmethod from_tensordict(tensordict, non_tensordict=None)¶

Tensor 類別包裝器，用於實例化一個新的 tensor 類別物件。

參數:

tensordict (TensorDict) – Tensor 類型的字典
non_tensordict (dict) – 包含非 tensor 和巢狀 tensor 類別物件的字典

get(key: NestedKey, default: Any = _NoDefault.ZERO)¶

取得與輸入鍵儲存的值。

參數:

key (str, tuple of str) – 要查詢的鍵。如果為 str 的 tuple，則等效於串聯呼叫 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¶

從磁碟載入記憶體映射 (memory-mapped) 的 tensordict。

參數:

prefix (str 或 資料夾路徑 (Path to folder)) – 應提取儲存的 tensordict 的資料夾路徑。
device (torch.device 或 等效物件 (equivalent), 選用 (optional)) – 如果提供，資料將會非同步地轉換 (cast) 到該裝置。支援 “meta” 裝置，在此情況下，不會載入資料，但會建立一組空的 "meta" 張量。這對於了解總模型大小和結構非常有用，而無需實際開啟任何檔案。
non_blocking (bool, 選用 (optional)) – 如果 True, 在將張量載入裝置後，將不會呼叫 synchronize。預設為 False。
out (TensorDictBase, 選用 (optional)) – 應寫入資料的選用 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" 裝置上，或者作為虛擬張量 (fake tensor)。

範例

>>> 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 就地 (in-place) 載入到目標 tensorclass。

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 中對應的記憶體映射張量 (memory-mapped Tensor)。

參數:

prefix (str) – 將儲存記憶體映射張量的目錄前綴 (directory prefix)。目錄樹結構將模仿 tensordict 的結構。
copy_existing (bool) – 如果為 False (預設)，如果 tensordict 中的條目 (entry) 已經是儲存在磁碟上且具有關聯檔案的張量，但未根據前綴儲存在正確的位置，則會引發例外。如果為 True，則任何現有張量都將複製到新位置。

關鍵字引數 (Keyword Arguments):

num_threads (int, 選用 (optional)) – 用於寫入 memmap 張量的執行緒數。預設為 0。
return_early (bool, 選用 (optional)) – 如果 True 且 num_threads>0，則該方法將傳回 tensordict 的 future。
share_non_tensor (bool, 選用 (optional)) – 如果為 True，則非張量資料將在進程和寫入操作之間共享（例如就地更新或設定），並且在單個節點內的任何 worker 上的操作將更新所有其他 worker 上的值。如果非張量葉子 (non-tensor leaves) 的數量很高（例如，共享大量的非張量資料堆疊），則可能會導致 OOM 或類似的錯誤。預設為 False。
existsok (bool, 選用 (optional)) – 如果為 False，如果張量已經存在於相同的路徑中，則會引發例外。預設為 True。

然後鎖定 TensorDict，這意味著任何非就地的寫入操作都會引發例外（例如，重新命名、設定或移除條目）。一旦 tensordict 被解鎖，memory-mapped 屬性將變為 False，因為跨進程身份 (cross-process identity) 不再保證。

傳回:: 如果 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¶

將所有 tensors 原地寫入對應的記憶體映射 Tensor。

參數:

prefix (str) – 將儲存記憶體映射張量的目錄前綴 (directory prefix)。目錄樹結構將模仿 tensordict 的結構。
copy_existing (bool) – 如果為 False (預設)，如果 tensordict 中的條目 (entry) 已經是儲存在磁碟上且具有關聯檔案的張量，但未根據前綴儲存在正確的位置，則會引發例外。如果為 True，則任何現有張量都將複製到新位置。

關鍵字引數 (Keyword Arguments):

num_threads (int, 選用 (optional)) – 用於寫入 memmap 張量的執行緒數。預設為 0。
return_early (bool, optional) – 如果 True 且 num_threads>0，此方法將會回傳 tensordict 的 future。可以使用 future.result() 查詢結果 tensordict。
share_non_tensor (bool, 選用 (optional)) – 如果為 True，則非張量資料將在進程和寫入操作之間共享（例如就地更新或設定），並且在單個節點內的任何 worker 上的操作將更新所有其他 worker 上的值。如果非張量葉子 (non-tensor leaves) 的數量很高（例如，共享大量的非張量資料堆疊），則可能會導致 OOM 或類似的錯誤。預設為 False。
existsok (bool, 選用 (optional)) – 如果為 False，如果張量已經存在於相同的路徑中，則會引發例外。預設為 True。

然後鎖定 TensorDict，這意味著任何非就地的寫入操作都會引發例外（例如，重新命名、設定或移除條目）。一旦 tensordict 被解鎖，memory-mapped 屬性將變為 False，因為跨進程身份 (cross-process identity) 不再保證。

傳回:: 如果 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 (str) – 將儲存記憶體映射張量的目錄前綴 (directory prefix)。目錄樹結構將模仿 tensordict 的結構。
copy_existing (bool) – 如果為 False (預設)，如果 tensordict 中的條目 (entry) 已經是儲存在磁碟上且具有關聯檔案的張量，但未根據前綴儲存在正確的位置，則會引發例外。如果為 True，則任何現有張量都將複製到新位置。

關鍵字引數 (Keyword Arguments):

num_threads (int, 選用 (optional)) – 用於寫入 memmap 張量的執行緒數。預設為 0。
return_early (bool, 選用 (optional)) – 如果 True 且 num_threads>0，則該方法將傳回 tensordict 的 future。
share_non_tensor (bool, 選用 (optional)) – 如果為 True，則非張量資料將在進程和寫入操作之間共享（例如就地更新或設定），並且在單個節點內的任何 worker 上的操作將更新所有其他 worker 上的值。如果非張量葉子 (non-tensor leaves) 的數量很高（例如，共享大量的非張量資料堆疊），則可能會導致 OOM 或類似的錯誤。預設為 False。
existsok (bool, 選用 (optional)) – 如果為 False，如果張量已經存在於相同的路徑中，則會引發例外。預設為 True。

然後鎖定 TensorDict，這意味著任何非就地的寫入操作都會引發例外（例如，重新命名、設定或移除條目）。一旦 tensordict 被解鎖，memory-mapped 屬性將變為 False，因為跨進程身份 (cross-process identity) 不再保證。

傳回:: 如果 return_early=False 則回傳一個資料儲存為記憶體映射 tensors 的新 TensorDict 實例，否則回傳 TensorDictFuture 實例。

注意

這是將一組大型 buffers 寫入磁碟的建議方法，因為 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，則重新整理其內容。

如果沒有與其相關聯的路徑，此方法將會引發例外。

property names: Size¶

檢索 tensor 類的 dim names。

傳回:: names (str 的 list)

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 的 tuple) – 要設定的鍵的名稱。如果是 str 的 tuple，則相當於鏈式呼叫 getattr，然後最後呼叫 setattr。
value (Any) – 要儲存在 tensorclass 中的值
inplace (bool, optional) – 如果 True，set 會嘗試就地更新值。如果 False 或鍵不存在，則該值將直接寫入其目的地。

傳回:

self

state_dict(destination=None, prefix='', keep_vars=False, flatten=False) → dict[str, Any]¶: 傳回一個 state_dict 字典，可用於儲存和載入 tensorclass 中的資料。

to_tensordict() → TensorDict¶

將 tensorclass 轉換為常規的 TensorDict。

建立所有條目的副本。 Memmap 和共享記憶體 tensors 會轉換為常規 tensors。

傳回:: 一個新的 TensorDict 物件，包含與 tensorclass 相同的值。

unbind(dim: int)¶

傳回一個索引 tensorclass 實例的 tuple，這些實例沿指示的維度解除綁定。

產生的 tensorclass 實例將共享初始 tensorclass 實例的儲存空間。

PromptData¶

文件

教學課程

資源