使用 TensorDict 简化 PyTorch 内存管理

作者： Tom Begley

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

设备

创建 时，可以使用 keyword argument 指定设备。如果设置了 ，则 的所有条目都将放置在该设备上。如果未设置 ，则 不要求 the 中的条目必须位于同一 装置。TensorDictdevicedeviceTensorDictdeviceTensorDict

在此示例中，我们使用 实例化 a 。什么时候 我们打印可以看到它们已移动到设备上的内容。TensorDictdevice="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)

如果 的设备不是 ，则还会移动新条目 拖动到设备上。TensorDictNone

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

您可以使用 该属性查看 的当前设备。TensorDictdevice

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

的内容可以发送到 PyTorch 张量等设备 with 或 with 成为所需的设备。TensorDictdevice

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

该方法需要一个有效的 device 作为参数传递。如果要从 to allow 值中删除 device 具有不同的设备，则应使用该方法。TensorDictTensorDict.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提供了一个类，它允许我们将张量的内容存储在磁盘上，同时仍然 支持快速索引和批量加载内容。 请参阅 ImageNet 教程以获取 实际示例。

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

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)

或者，可以使用该方法。这将 创建一个具有相同 values 的相同结构的新结构，但它不会复制 原始张量的内容添加到 内存映射张量。这允许您创建内存映射，然后缓慢填充它，因此通常应该是 首选 。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.])

默认情况下，它的内容将保存到一个临时的 位置，但是，如果您想控制它们的保存位置，您可以 使用 keyword 参数 。TensorDictprefix="/path/to/root"

的内容保存在一个目录结构中，该目录结构模拟 本身的结构。张量的内容被保存 在 NumPy memmap 中，以及关联的 PyTorch 保存文件中的元数据。例如 以上内容保存如下：TensorDictTensorDictTensorDict

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