简化PyTorch内存管理的TensorDict

作者: Tom Begley

在这个教程中，您将学习如何控制TensorDict的内容存储在内存中的位置，无论是通过将这些内容发送到设备，还是通过利用内存映射。

设备

当你创建一个TensorDict时，你可以通过device关键字参数指定一个设备。如果device被设置，则TensorDict的所有条目都将放置在该设备上。如果device未设置，则TensorDict中的条目不需要在同一设备上。

在这个示例中，我们实例化了一个 TensorDict 并使用 device="cuda:0"。当我们打印内容时，可以看到它们已经被移动到了设备上。

>>> import torch
>>> from tensordict import TensorDict
>>> tensordict = TensorDict({"a": torch.rand(10)}, [10], device="cuda:0")
>>> print(tensordict)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([10]), device=cuda:0, dtype=torch.float32, is_shared=True)},
    batch_size=torch.Size([10]),
    device=cuda:0,
    is_shared=True)

如果设备的 TensorDict 不是 None，新条目也会被移动到该设备上。

>>> tensordict["b"] = torch.rand(10, 10)
>>> print(tensordict)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([10]), device=cuda:0, dtype=torch.float32, is_shared=True),
        b: Tensor(shape=torch.Size([10, 10]), device=cuda:0, dtype=torch.float32, is_shared=True)},
    batch_size=torch.Size([10]),
    device=cuda:0,
    is_shared=True)

您可以使用 TensorDict 的 device 属性来检查当前设备。

>>> print(tensordict.device)
cuda:0

TensorDict 的内容可以发送到设备，如 PyTorch 张量 使用 TensorDict.cuda() 或 TensorDict.device(device) 其中 device 是目标设备。

>>> tensordict.to(torch.device("cpu"))
>>> print(tensordict)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([10, 10]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([10]),
    device=cpu,
    is_shared=False)
>>> tensordict.cuda()
>>> print(tensordict)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([10]), device=cuda:0, dtype=torch.float32, is_shared=True),
        b: Tensor(shape=torch.Size([10, 10]), device=cuda:0, dtype=torch.float32, is_shared=True)},
    batch_size=torch.Size([10]),
    device=cuda:0,
    is_shared=True)

The TensorDict.device 方法需要传递一个有效的设备作为参数。如果你想从 TensorDict 中移除设备以允许具有不同设备的值，你应该使用 TensorDict.clear_device 方法。

>>> tensordict.clear_device()
>>> print(tensordict)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([10]), device=cuda:0, dtype=torch.float32, is_shared=True),
        b: Tensor(shape=torch.Size([10, 10]), device=cuda:0, dtype=torch.float32, is_shared=True)},
    batch_size=torch.Size([10]),
    device=None,
    is_shared=False)

内存映射张量

tensordict 提供了一个类 MemoryMappedTensor 它允许我们将张量的内容存储在磁盘上，同时仍然 支持快速索引和批量加载内容。 请参阅ImageNet 教程以了解实际应用示例。

要将TensorDict转换为内存映射张量的集合，请使用 TensorDict.memmap_。

tensordict = TensorDict({"a": torch.rand(10), "b": {"c": torch.rand(10)}}, [10])
tensordict.memmap_()

print(tensordict)

TensorDict(
    fields={
        a: MemoryMappedTensor(shape=torch.Size([10]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: MemoryMappedTensor(shape=torch.Size([10]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([10]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([10]),
    device=cpu,
    is_shared=False)

或者可以使用 TensorDict.memmap_like 方法。这将创建一个新的 TensorDict，具有相同的结构和 MemoryMappedTensor 值，但不会将原始张量的内容复制到内存映射张量中。这允许你创建内存映射的 TensorDict，然后慢慢填充它，因此通常应优先于 memmap_。

tensordict = TensorDict({"a": torch.rand(10), "b": {"c": torch.rand(10)}}, [10])
mm_tensordict = tensordict.memmap_like()

print(mm_tensordict["a"].contiguous())

MemoryMappedTensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])

默认情况下，TensorDict 的内容将被保存到磁盘上的临时位置，但是如果您希望控制它们的保存位置，可以使用关键字参数 prefix="/path/to/root"。

TensorDict 的内容保存在一个目录结构中，该结构模仿了 TensorDict 本身的结构。张量的内容以 NumPy memmap 的形式保存，元数据则保存在关联的 PyTorch save 文件中。例如，上述 TensorDict 保存如下：

├── a.memmap
├── a.meta.pt
├── b
│ ├── c.memmap
│ ├── c.meta.pt
│ └── meta.pt
└── meta.pt