PairwiseDataset¶

class torchrl.data.PairwiseDataset(chosen_data: 'RewardData', rejected_data: 'RewardData', *, batch_size, device=None, names=None)[原始碼]¶

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

dumps(prefix: str | None = None, copy_existing: bool =False, *, num_threads: int =0, return_early: bool =False, share_non_tensor: bool =False) → T¶

將 tensordict 儲存到磁碟。

此函式是 memmap() 的代理。

classmethod fields()¶

返回描述此 dataclass 欄位的元組。

接受一個 dataclass 或其例項。元組元素型別為 Field。

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

根據資料集名稱返回一個 PairwiseDataset 例項。

引數：:

split (字串) – "train" 或 "valid"，取決於所需的資料分割。
dataset_name (字串, 可選) – 待處理資料集的名稱。預設為 "CarperAI/openai_summarize_comparisons"。
max_length (整數, 可選) – 資料集序列的最大長度。預設為 550。
root_dir (路徑, 可選) – 資料集儲存的路徑。預設為 "$HOME/.cache/torchrl/data"
from_disk (布林值, 可選) – 如果為 True，將使用 datasets.load_from_disk()。否則，將使用 datasets.load_dataset()。預設為 False。

返回：包含記憶體對映版本的所需資料集的 PairwiseDataset 例項，: 所需資料集的版本。

示例

>>> data = PairwiseDataset.from_dataset("train")
>>> print(data)
PairwiseDataset(
    chosen_data=RewardData(
        attention_mask=MemoryMappedTensor(shape=torch.Size([92534, 550]), device=cpu, dtype=torch.int64, is_shared=False),
        input_ids=MemoryMappedTensor(shape=torch.Size([92534, 550]), device=cpu, dtype=torch.int64, is_shared=False),
        rewards=None,
        end_scores=None,
        batch_size=torch.Size([92534]),
        device=None,
        is_shared=False),
    rejected_data=RewardData(
        attention_mask=MemoryMappedTensor(shape=torch.Size([92534, 550]), device=cpu, dtype=torch.int64, is_shared=False),
        input_ids=MemoryMappedTensor(shape=torch.Size([92534, 550]), device=cpu, dtype=torch.int64, is_shared=False),
        rewards=None,
        end_scores=None,
        batch_size=torch.Size([92534]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([92534]),
    device=None,
    is_shared=False)
>>> # data can be sampled from using regular indexing
>>> sub_data = data[:3]

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

張量類包裝器，用於例項化新的張量類物件。

引數：:

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

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

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

引數：:

key (字串, 字串元組) – 要查詢的鍵。如果是字串元組，則相當於鏈式呼叫 getattr。
default – 如果在 tensorclass 中找不到鍵，則返回的預設值。

返回：:

使用輸入鍵儲存的值

classmethod load(prefix: str | Path, *args, **kwargs) → T¶

從磁碟載入 tensordict。

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

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

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

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

classmethod load_memmap(prefix: str | Path, device: torch.device | None = None, non_blocking: bool =False, *, out: TensorDictBase | None = 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()

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

示例

>>> 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)¶: 嘗試在目標 tensorclass 上就地載入 state_dict。

memmap(prefix: str | None = 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 (布林值, 可選) – 如果為 True 且 num_threads>0，方法將返回一個 tensordict 的 future。
share_non_tensor (布林值, 可選) – 如果為 True，非張量資料將在程序之間共享，並且在單個節點內任何工作程序上的寫入操作（例如就地更新或設定）將更新所有其他工作程序上的值。如果非張量葉子數量很高（例如，共享大量的非張量資料堆疊），這可能會導致 OOM 或類似錯誤。預設為 False。
existsok (布林值, 可選) – 如果為 False，如果同一路徑中已存在張量，則會引發異常。預設為 True。

該 TensorDict 將被鎖定，這意味著任何非就地的寫入操作（例如，重新命名、設定或移除條目）都將丟擲異常。一旦 tensordict 被解鎖，memory-mapped 屬性將設定為 False，因為跨程序身份不再得到保證。

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

注意

以這種方式序列化深度巢狀的 tensordicts 可能會很慢，因此不建議在訓練迴圈內呼叫此方法。

memmap_(prefix: str | None = 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 (布林值, 可選) – 如果為 True 且 num_threads>0，方法將返回一個 tensordict 的 future。可以使用 future.result() 查詢結果 tensordict。
share_non_tensor (布林值, 可選) – 如果為 True，非張量資料將在程序之間共享，並且在單個節點內任何工作程序上的寫入操作（例如就地更新或設定）將更新所有其他工作程序上的值。如果非張量葉子數量很高（例如，共享大量的非張量資料堆疊），這可能會導致 OOM 或類似錯誤。預設為 False。
existsok (布林值, 可選) – 如果為 False，如果同一路徑中已存在張量，則會引發異常。預設為 True。

該 TensorDict 將被鎖定，這意味著任何非就地的寫入操作（例如，重新命名、設定或移除條目）都將丟擲異常。一旦 tensordict 被解鎖，memory-mapped 屬性將設定為 False，因為跨程序身份不再得到保證。

返回：:: 如果 return_early=False，則返回 self；否則，返回一個 TensorDictFuture 例項。

注意

以這種方式序列化深度巢狀的 tensordicts 可能會很慢，因此不建議在訓練迴圈內呼叫此方法。

memmap_like(prefix: str | None = 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 (布林值, 可選) – 如果為 True 且 num_threads>0，方法將返回一個 tensordict 的 future。
share_non_tensor (布林值, 可選) – 如果為 True，非張量資料將在程序之間共享，並且在單個節點內任何工作程序上的寫入操作（例如就地更新或設定）將更新所有其他工作程序上的值。如果非張量葉子數量很高（例如，共享大量的非張量資料堆疊），這可能會導致 OOM 或類似錯誤。預設為 False。
existsok (布林值, 可選) – 如果為 False，如果同一路徑中已存在張量，則會引發異常。預設為 True。

該 TensorDict 將被鎖定，這意味著任何非就地的寫入操作（例如，重新命名、設定或移除條目）都將丟擲異常。一旦 tensordict 被解鎖，memory-mapped 屬性將設定為 False，因為跨程序身份不再得到保證。

返回：:: 如果 return_early=False，則返回一個數據以記憶體對映張量形式儲存的新 TensorDict 例項；否則，返回一個 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: str | None = 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 (字串, 字串元組) – 要設定的鍵的名稱。如果是字串元組，則相當於鏈式呼叫 getattr，最後呼叫 setattr。
value (Any) – 要儲存在 tensorclass 中的值
inplace (布林值, 可選) – 如果為 True，set 將嘗試就地更新值。如果為 False 或鍵不存在，值將簡單地寫入其目的地。

返回：:

self

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

to_tensordict(*, retain_none: bool | None = None) → TensorDict¶

將張量類轉換為常規張量字典。

複製所有條目。記憶體對映張量和共享記憶體張量將轉換為常規張量。

引數：:

retain_none (bool) –

如果為 True，則 None 值將被寫入張量字典。否則它們將被丟棄。預設值：True。

注意

從 v0.8 版本開始，預設值將切換為 False。

返回：:

一個新的張量字典物件，包含與張量類相同的值。

unbind(dim: int)¶

返回一個元組，其中包含沿指定維度分解的帶索引的張量類例項。

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

PairwiseDataset¶

文件

教程

資源