使用 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.cuda() 或 TensorDict.device(device) 将 TensorDict 的内容发送到设备，类似于 PyTorch 张量，其中 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 方法。这将创建一个具有 MemoryMappedTensor 值的新 TensorDict，其结构相同，但不会将原始张量的内容复制到内存映射张量。这使您可以创建内存映射 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