PyTorch中如何将dstack堆叠的张量转为按列堆叠的hstack结果

直接使用transpose交换维度后展平即可，单行实现代码如下：

h_from_d = d.transpose(1, 2).flatten(start_dim=1)

如果堆叠的张量数量固定，也可以直接拆分拼接实现，效果完全一致：

h_from_d = torch.cat([d[..., 0], d[..., 1]], dim=1)

原理说明 #

• 原始a、b形状均为(4, 2)，经过dstack后得到的d形状为(4, 2, 2)，三个维度分别对应：行号、原张量的列号、堆叠的张量序号（0对应a，1对应b）
• 直接reshape/flatten会按内存存储顺序读取元素，导致顺序混乱，因此需要先交换后两个维度，将「堆叠的张量序号」维度调整到「原张量的列号」维度之前，调整后每个行对应的二维切片结构为[[a的该行元素], [b的该行元素]]
• 最后从第1维开始展平，即可得到和hstack完全一致的结果

验证代码 #

你可以运行以下代码确认结果正确性：

import torch

a = torch.tensor(([1,2],
                  [3,4],
                  [5,6],
                  [7,8]))
b = torch.tensor(([0,0],
                  [1,1],
                  [2,2],
                  [3,3]))

d = torch.dstack([a, b])
h = torch.hstack([a, b])

h_from_d = d.transpose(1, 2).flatten(start_dim=1)
print(torch.allclose(h, h_from_d)) # 输出True，说明结果完全一致

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