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)[source]¶
返回 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()
的别名。