使用 TensorDict 预分配内存

作者: Tom Begley

在本教程中，您将学习如何利用 TensorDict 中的内存预分配。

假设我们有一个函数，它返回一个 TensorDict

import torch
from tensordict.tensordict import TensorDict


def make_tensordict():
    return TensorDict({"a": torch.rand(3), "b": torch.rand(3, 4)}, [3])

也许我们想多次调用此函数，并使用结果填充单个 TensorDict。

N = 10
tensordict = TensorDict({}, batch_size=[N, 3])

for i in range(N):
    tensordict[i] = make_tensordict()

print(tensordict)

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

因为我们已经指定了 tensordict 的 batch_size，所以在循环的第一次迭代期间，我们使用空张量填充 tensordict，这些空张量的第一个维度大小为 N，其余维度由 make_tensordict 的返回值确定。 在上面的示例中，我们为键 "a" 预分配了大小为 torch.Size([10, 3]) 的零数组，为键 "b" 预分配了大小为 torch.Size([10, 3, 4]) 的数组。 循环的后续迭代将就地写入。 因此，如果并非所有值都被填充，它们将获得默认值零。

让我们通过逐步执行上面的循环来演示正在发生的事情。 我们首先初始化一个空的 TensorDict。

N = 10
tensordict = TensorDict({}, batch_size=[N, 3])
print(tensordict)

TensorDict(
    fields={
    },
    batch_size=torch.Size([10, 3]),
    device=None,
    is_shared=False)

在第一次迭代之后，tensordict 已预先填充了 "a" 和 "b" 的张量。 这些张量包含零，但我们已为其分配随机值的第一行除外。

random_tensordict = make_tensordict()
tensordict[0] = random_tensordict

assert (tensordict[1:] == 0).all()
assert (tensordict[0] == random_tensordict).all()

print(tensordict)

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

在后续迭代中，我们就地更新预分配的张量。

a = tensordict["a"]
random_tensordict = make_tensordict()
tensordict[1] = random_tensordict

# the same tensor is stored under "a", but the values have been updated
assert tensordict["a"] is a
assert (tensordict[:2] != 0).all()