PyTorch 如何按自定义顺序高效实现张量重塑操作？

你需要的重排不属于普通的维度变换，本质是先对原张量做分块、调整分块顺序后再合并维度，下面是两种最高效的实现方案：

方法一：内置维度变换（推荐，零额外内存开销） #

所有操作均为PyTorch视图操作，不会产生数据拷贝，性能最优：

import torch
t = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14 ,15, 16, 17])
# 步骤1：把1维张量拆分为 3个块 × 2行 × 3列 的结构
t_reshaped = t.reshape(3, 2, 3)
# 步骤2：交换维度，把「行」维度挪到最前，「块」维度放到中间
t_permuted = t_reshaped.permute(1, 0, 2)
# 步骤3：合并后面两个维度得到最终结果
t_reshape = t_permuted.flatten(start_dim=1)

也可以把步骤合并成一行：

t_reshape = t.reshape(3,2,3).permute(1,0,2).flatten(1)

方法二：索引映射（适合更复杂的自定义重排规则） #

如果后续重排逻辑会调整得更复杂，可以直接构造目标位置对应的原索引数组，直接索引取值即可：

# 构造索引映射：目标张量顺序对应的原张量下标
idx = torch.tensor([0,1,2,6,7,8,12,13,14,3,4,5,9,10,11,15,16,17])
t_reshape = t[idx].reshape(2, 9)

这个方法灵活性最高，只要提前算好索引规则就能适配任意重排需求，性能也足够应对绝大多数场景。

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