PyTorch——TensorDict类的使用

整体说明

  • TensorDict 的本质是：一个统一管理多个张量的容器，并强制规定它们共享同一个“批次结构（Batch Shape）”
    • 普通字典 ：只管存数据，维度是否对齐全靠用户记忆整理
    • TensorDict ：通过 batch_size 定义“骨架”，确保所有存进去的张量在批次维度上严格对齐

TensorDict 相关示例

  • TensorDict 相关示例

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    from tensordict import TensorDict
    import torch

    # 标准创建：batch_size=[3,4] 代表一个 3行4列 的二维批次网格
    td = TensorDict(
        {
            "obs": torch.randn(3, 4, 10),   # 每个网格点存一个10维观测向量
            "action": torch.randn(3, 4, 2)  # 每个网格点存一个2维动作向量
        },
        batch_size=[3, 4]
    )
    • 所有张量的前 len(batch_size) 个维度（即前 2 维）必须严格等于 [3, 4] ，剩下的维度（如 10、2）被视为各自独立的“特征维度”，可以各不相同
    • 后续添加张量是也需要满足 前 2 维 必须严格等于 [3, 4]

TensorDict 核心操作

  • 核心操作示例 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    85
    86
    87
    88
    89
    90
    91
    92
    from tensordict import TensorDict
    import torch

    # 标准创建：batch_size=[3,4] 代表一个 3行4列 的二维批次网格
    td = TensorDict(
        {
            "obs": torch.randn(3, 4, 10),   # 每个网格点存一个10维观测向量
            "action": torch.randn(3, 4, 2)  # 每个网格点存一个2维动作向量
        },
        batch_size=[3, 4]
    )
    print(td)
    # TensorDict(
    #     fields={
    #         action: Tensor(shape=torch.Size([3, 4, 2]), device=cpu, dtype=torch.float32, is_shared=False),
    #         obs: Tensor(shape=torch.Size([3, 4, 10]), device=cpu, dtype=torch.float32, is_shared=False)},
    #     batch_size=torch.Size([3, 4]),
    #     device=None,
    #     is_shared=False)

    # 1 像字典一样存取键值
    td["reward"] = torch.randn(3, 4, 1)   # 新增键
    print(td)

    obs_data = td["obs"]                  # 获取值，支持 get()
    print(obs_data) # 获取 obs 对应的 tensor 张量（普通张量）

    # 2 像张量一样索引（基于 batch_size 切片）
    # 因为系统知道“前2维是批次”，所以你可以直接按网格索引，它会自动对所有键同步切片：
    sub_td = td[1, 2]          # 取出第1行第2列，返回新TensorDict
    # 此时 sub_td["obs"] 自动变成 (10,)，sub_td["action"] 自动变成 (2,)
    print(sub_td)
    # TensorDict(
    #     fields={
    #         action: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.float32, is_shared=False),
    #         obs: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.float32, is_shared=False),
    #         reward: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False)},
    #     batch_size=torch.Size([]),
    #     device=None,
    #     is_shared=False)

    td1 = td[1] # 取出第 1 行，返回新 TensorDict
    print(td1)
    # TensorDict(
    #     fields={
    #         action: Tensor(shape=torch.Size([4, 2]), device=cpu, dtype=torch.float32, is_shared=False),
    #         obs: Tensor(shape=torch.Size([4, 10]), device=cpu, dtype=torch.float32, is_shared=False),
    #         reward: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
    #     batch_size=torch.Size([4]),
    #     device=None,
    #     is_shared=False)

    batch_td = td[:, 0]        # 取出所有行的第0列，形状变为 [3, ...]
    print(batch_td)
    # TensorDict(
    #     fields={
    #         action: Tensor(shape=torch.Size([3, 2]), device=cpu, dtype=torch.float32, is_shared=False),
    #         obs: Tensor(shape=torch.Size([3, 10]), device=cpu, dtype=torch.float32, is_shared=False),
    #         reward: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
    #     batch_size=torch.Size([3]),
    #     device=None,
    #     is_shared=False)

    # 3 像张量一样重塑形状（自动保护特征维）
    # 将 3x4 的网格展平为 12 个独立样本
    flat_td = td.reshape(12)
    # flat_td["obs"] 自动变成 (12, 10)，特征维度 10 纹丝不动
    print(flat_td)
    # TensorDict(
    #     fields={
    #         action: Tensor(shape=torch.Size([12, 2]), device=cpu, dtype=torch.float32, is_shared=False),
    #         obs: Tensor(shape=torch.Size([12, 10]), device=cpu, dtype=torch.float32, is_shared=False),
    #         reward: Tensor(shape=torch.Size([12, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
    #     batch_size=torch.Size([12]),
    #     device=None,
    #     is_shared=False)

    # 变换为 2x6 的网格
    new_td = td.reshape(2, 6)
    print(new_td)
    # TensorDict(
    #     fields={
    #         action: Tensor(shape=torch.Size([2, 6, 2]), device=cpu, dtype=torch.float32, is_shared=False),
    #         obs: Tensor(shape=torch.Size([2, 6, 10]), device=cpu, dtype=torch.float32, is_shared=False),
    #         reward: Tensor(shape=torch.Size([2, 6, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
    #     batch_size=torch.Size([2, 6]),
    #     device=None,
    #     is_shared=False)


    # 4 设备管理与转移
    td = td.to("cpu")       # 所有内部张量一键转移至 CPU

补充：高级用法

  • 嵌套结构 ：支持嵌套字典或元组键，方便管理复杂数据（如多智能体）

    1
    td["agent", "position"] = torch.randn(3, 4, 2)

  • 与 PyTorch 模块无缝协作 ：可以编写一个接受 TensorDict 的神经网络，模型内部只需关心特征维度（如 x = td["obs"]），完全不用处理 batch_size 的变换，极大提高代码复用性

  • 批量操作（Stack / Cat） ：拼接多个 TensorDict 时，系统会严格校验批次维度是否兼容，避免人工对齐错误

  • 内存映射（Memory-mapping） ：支持将超大张量映射到磁盘，无需全部加载到内存即可操作，适合超大规模数据集

理解：为什么一定要定义 batch_size？

  • batch_size 告诉 TensorDict 哪些维度是 Batch 维度，哪些维度是 特征维度
    • 当执行 td[0]、td.reshape() 或 td.stack() 时，系统依赖这把标尺自动推断如何修改所有内部张量
    • 有了 batch_size 后，再也不需要手动写 obs = d["obs"][batch_idx] 和 act = d["action"][batch_idx]，一次索引全部搞定