如何将平铺的PyTorch张量拆分为嵌套小张量列表？

步骤1：正确重塑张量形状 #

你的原始张量形状是 (24, 32)，要拆分为6个外层元素、每个元素含8个4×4张量的结构，首先需要正确调整张量维度：

• 24行可拆分为6组，每组4行（6×4=24）
• 32列可拆分为8组，每组4列（8×4=32）

方法一：reshape + permute #

先将张量转为 (6,4,8,4) 形状，再交换中间两个维度得到 (6,8,4,4)，对应6×8的4×4张量网格：

import torch

original_tensor = torch.randn(24, 32)  # 示例输入张量
# 按行分组为6个4行块，按列分组为8个4列块，再调整维度顺序
tensor_to_split = original_tensor.reshape(6, 4, 8, 4).permute(0, 2, 1, 3)
print(tensor_to_split.shape)  # 输出: torch.Size([6, 8, 4, 4])

方法二：unfold直接提取块 #

用unfold方法更直观地提取指定大小的子块：

# 沿行维度提取步长为4的4行块，沿列维度提取步长为4的4列块
tensor_to_split = original_tensor.unfold(0, 4, 4).unfold(1, 4, 4)
print(tensor_to_split.shape)  # 输出: torch.Size([6, 8, 4, 4])

步骤2：转为Pythonic的嵌套列表 #

得到 (6,8,4,4) 形状的张量后，无需复杂嵌套循环，用列表推导式即可快速转为目标结构：

desired_tensors = [list(block_row) for block_row in tensor_to_split]

• 遍历tensor_to_split的第一维度（6个元素），每个元素是(8,4,4)形状的张量
• 将每个(8,4,4)张量转为列表，得到8个4×4张量的子列表
• 最终desired_tensors是包含6个子列表的列表，每个子列表含8个4×4张量

验证结构正确性：

print(len(desired_tensors))       # 输出: 6
print(len(desired_tensors[0]))    # 输出: 8
print(desired_tensors[0][0].shape)# 输出: torch.Size([4,4])

替代方案：显式循环写法（若偏好直观逻辑） #

如果更习惯显式循环，可简化你原始代码的语法错误后使用：

desired_tensors = []
for r in range(6):
    clist = []
    for c in range(8):
        clist.append(tensor_to_split[r, c])  # 逗号索引等价于tensor_to_split[r][c]，更高效
    desired_tensors.append(clist)

内容的提问来源于stack exchange，提问作者Christopher Clark