使用 TensorDict 简化 PyTorch 内存管理

作者: Tom Begley

在本教程中，您将学习如何控制 TensorDict 的内容在内存中的存储位置，可以通过将这些内容发送到设备，或者利用内存映射。

设备

创建 TensorDict 时，可以使用 device 关键字参数指定设备。如果设置了 device，则 TensorDict 的所有条目都将放置在该设备上。如果未设置 device，则 TensorDict 中的条目不必在同一设备上。

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

>>> 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)

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

>>> 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)

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 内存映射文件中，元数据保存在相关的 PyTorch 保存文件中。例如，上面的 TensorDict 会保存如下：

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