快捷方式

LazyStackStorage

class torchrl.data.replay_buffers.LazyStackStorage(max_size: int | None = None, *, compilable: bool = False, stack_dim: int = 0, device: torch.device | str | int | None = None)[源代码]

返回 LazyStackTensorDict 实例的 ListStorage。

此存储允许将异构结构索引为单个 TensorDict 表示。它使用 LazyStackedTensorDict,它操作非连续的 tensordicts 列表,并在查询时惰性地堆叠项。这意味着此存储的采样速度会很快,但数据访问可能会很慢(因为它需要堆叠)。存储中还可以存储具有异构形状的张量并一起堆叠。由于存储被表示为一个列表,因此需要存储在内存中的张量数量将随着缓冲区的大小呈线性增长。

如果可能,还可以通过 densify() 创建嵌套张量(请参阅 nested)。

参数:

max_size (int, optional) – 存储中存储的最大元素数量。如果未提供,则创建一个无限制的存储。

关键字参数:
  • compilable (bool, optional) – 如果为 True,则存储将与 compile() 兼容,但代价是在多进程环境中执行。

  • stack_dim (int, optional) – TensorDict 批次大小方面的堆叠维度。默认为 0

  • device (str, optional) – 要用于存储的设备。默认为 None(输入不会移动到设备)。

示例

>>> import torch
>>> from torchrl.data import ReplayBuffer, LazyStackStorage
>>> from tensordict import TensorDict
>>> _ = torch.manual_seed(0)
>>> rb = ReplayBuffer(storage=LazyStackStorage(max_size=1000, stack_dim=-1))
>>> data0 = TensorDict(a=torch.randn((10,)), b=torch.rand(4), c="a string!")
>>> data1 = TensorDict(a=torch.randn((11,)), b=torch.rand(4), c="another string!")
>>> _ = rb.add(data0)
>>> _ = rb.add(data1)
>>> rb.sample(10)
LazyStackedTensorDict(
    fields={
        a: Tensor(shape=torch.Size([10, -1]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([10, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        c: NonTensorStack(
            ['another string!', 'another string!', 'another st...,
            batch_size=torch.Size([10]),
            device=None)},
    exclusive_fields={
    },
    batch_size=torch.Size([10]),
    device=None,
    is_shared=False,
    stack_dim=0)
attach(buffer: Any) None

此函数将采样器附加到此存储。

从该存储读取的缓冲区必须通过调用此方法作为已附加实体包含进来。这确保了当存储中的数据发生变化时,组件能够感知到这些变化,即使该存储与其他缓冲区(例如,Priority Samplers)共享。

参数:

buffer – 读取此存储的对象。

dump(*args, **kwargs)

dumps() 的别名。

load(*args, **kwargs)

loads() 的别名。

save(*args, **kwargs)

dumps() 的别名。

文档

访问全面的 PyTorch 开发者文档

查看文档

教程

为初学者和高级开发者提供深入的教程

查看教程

资源

查找开发资源并让您的问题得到解答

查看资源