TensorStorage

class torchrl.data.replay_buffers.TensorStorage(storage, max_size=None, *, device: device = 'cpu', ndim: int = 1)

张量和张量字典的存储。

参数：

• storage (tensor 或 TensorDict) – 要使用的的数 据缓冲区。

• max_size (int) – 存储的大小，即缓冲区中存储的元素的最大数量。

关键字参数：

• device (torch.device, 可选) – 采样张量将存储和发送到的设备。默认为 torch.device("cpu")。如果传递“auto”，则会自动从传递的第一批数据中收集设备。默认情况下不启用此功能，以避免错误地将数据放置在 GPU 上，从而导致内存不足问题。

• ndim (int, 可选) – 用于测量存储大小时要考虑的维度数。例如，形状为 [3, 4] 的存储如果 ndim=1 则容量为 3，如果 ndim=2 则容量为 12。默认为 1。

示例

>>> data = TensorDict({
...     "some data": torch.randn(10, 11),
...     ("some", "nested", "data"): torch.randn(10, 11, 12),
... }, batch_size=[10, 11])
>>> storage = TensorStorage(data)
>>> len(storage)  # only the first dimension is considered as indexable
10
>>> storage.get(0)
TensorDict(
    fields={
        some data: Tensor(shape=torch.Size([11]), device=cpu, dtype=torch.float32, is_shared=False),
        some: TensorDict(
            fields={
                nested: TensorDict(
                    fields={
                        data: Tensor(shape=torch.Size([11, 12]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([11]),
                    device=None,
                    is_shared=False)},
            batch_size=torch.Size([11]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([11]),
    device=None,
    is_shared=False)
>>> storage.set(0, storage.get(0).zero_()) # zeros the data along index ``0``

此类也支持 tensorclass 数据。

示例

>>> from tensordict import tensorclass
>>> @tensorclass
... class MyClass:
...     foo: torch.Tensor
...     bar: torch.Tensor
>>> data = MyClass(foo=torch.randn(10, 11), bar=torch.randn(10, 11, 12), batch_size=[10, 11])
>>> storage = TensorStorage(data)
>>> storage.get(0)
MyClass(
    bar=Tensor(shape=torch.Size([11, 12]), device=cpu, dtype=torch.float32, is_shared=False),
    foo=Tensor(shape=torch.Size([11]), device=cpu, dtype=torch.float32, is_shared=False),
    batch_size=torch.Size([11]),
    device=None,
    is_shared=False)

attach(buffer: Any) → None

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

从此存储读取的缓冲区必须通过调用此方法作为附加实体包含在内。这保证了当存储中的数据更改时，即使存储与其他缓冲区共享（例如优先采样器），组件也会意识到更改。

参数：

buffer – 从此存储读取的对象。

dump(*args, **kwargs)

dumps() 的别名。

load(*args, **kwargs)

loads() 的别名。

save(*args, **kwargs)

dumps() 的别名。