RewardData¶

class torchrl.data.RewardData(input_ids: 'torch.Tensor', attention_mask: 'torch.Tensor', rewards: 'Optional[torch.Tensor]' = None, end_scores: 'Optional[torch.Tensor]' = None, *, batch_size, device=None, names=None)[原始碼]¶

property batch_size: Size¶

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

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

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

classmethod from_tensordict(tensordict, non_tensordict=None)¶

張量類別封裝器，用於實例化新的張量類別物件。

參數:

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

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

取得與輸入鍵值相關聯儲存的值。

參數:

key (str, str 的元組) – 要查詢的鍵值。如果是 str 的元組，則相當於鏈式調用 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 (str 或 資料夾的路徑) – 應提取儲存的 tensordict 的資料夾路徑。
device (torch.device 或 等效裝置, 選用) – 如果提供，資料將非同步地轉換為該裝置。支援 “meta” 裝置，在這種情況下，不會載入資料，而是建立一組空的 “meta” 張量。這對於了解模型的總大小和結構而無需實際開啟任何檔案很有用。
non_blocking (bool, 選用) – 如果 True，在將張量載入裝置後，將不會調用同步。預設值為 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)¶: 嘗試在目標 tensorclass 上就地載入 state_dict。

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

關鍵字參數:

num_threads (int, optional) – 用於寫入記憶體映射張量的執行緒數。預設值為 0。
return_early (bool, optional) – 如果 True 且 num_threads>0，該方法將返回 tensordict 的 future。
share_non_tensor (bool, optional) – 如果 True，非張量資料將在進程和寫入操作之間共享（例如，就地更新或設定），並且單個節點內任何 worker 上的操作都將更新所有其他 worker 上的值。如果非張量葉節點的數量很大（例如，共享大量的非張量資料堆疊），這可能會導致記憶體不足 (OOM) 或類似的錯誤。預設值為 False。
existsok (bool, optional) – 如果 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 (str) – 記憶體映射張量將儲存的目錄前綴。目錄樹結構將模仿 tensordict 的結構。
copy_existing (bool) – 如果為 False（預設值），則當 tensordict 中的條目已經是儲存在磁碟上的張量，並且具有關聯的檔案，但根據前綴沒有儲存在正確的位置時，將會引發例外。如果 True，任何現有的張量將被複製到新的位置。

關鍵字參數:

num_threads (int, optional) – 用於寫入記憶體映射張量的執行緒數。預設值為 0。
return_early (bool, optional) – 如果 True 且 num_threads>0，該方法將返回 tensordict 的 future。可以使用 future.result() 查詢結果 tensordict。
share_non_tensor (bool, optional) – 如果 True，非張量資料將在進程和寫入操作之間共享（例如，就地更新或設定），並且單個節點內任何 worker 上的操作都將更新所有其他 worker 上的值。如果非張量葉節點的數量很大（例如，共享大量的非張量資料堆疊），這可能會導致記憶體不足 (OOM) 或類似的錯誤。預設值為 False。
existsok (bool, optional) – 如果 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 (str) – 記憶體映射張量將儲存的目錄前綴。目錄樹結構將模仿 tensordict 的結構。
copy_existing (bool) – 如果為 False（預設值），則當 tensordict 中的條目已經是儲存在磁碟上的張量，並且具有關聯的檔案，但根據前綴沒有儲存在正確的位置時，將會引發例外。如果 True，任何現有的張量將被複製到新的位置。

關鍵字參數:

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

然後 TensorDict 會被鎖定，這表示任何非就地的寫入操作都會拋出例外（例如，重新命名、設定或移除條目）。一旦 tensordict 被解鎖，記憶體映射屬性就會被設為 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，則刷新其內容。

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

property names: Size¶

檢索張量類別的維度名稱。

傳回:: names (字串列表)

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, 字串元組) – 要設定的鍵的名稱。如果為字串元組，則等同於鏈式呼叫 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 和共享記憶體張量會轉換為常規張量。

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

unbind(dim: int)¶

傳回沿指示維度取消綁定的索引 tensorclass 實例的元組。

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

RewardData¶

文件

教學課程

資源