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¶

文档

教程

资源