如何交错拼接两个PyTorch tensor矩阵以得到指定结构的输出

实现方案 #

方案1：更简洁的堆叠重排实现（推荐） #

直接通过维度堆叠+重排完成交错，不需要额外补零运算，性能更高：

import torch

# 定义原始矩阵
A = torch.tensor([[1,1,1,1,1],
                  [1,1,1,1,1],
                  [1,1,1,1,1]])

B = torch.tensor([[2,2,2,2,2],
                  [2,2,2,2,2],
                  [2,2,2,2,2]])

# 沿最后一维堆叠A、B，再将最后两维展平完成交错
C = torch.stack([A, B], dim=-1).flatten(start_dim=-2)

# 注：你示例中的C每行只有8个元素属于笔误，5个A元素和5个B元素交错后每行应为10个元素，输出结果为：
# tensor([[1, 2, 1, 2, 1, 2, 1, 2, 1, 2],
#         [1, 2, 1, 2, 1, 2, 1, 2, 1, 2],
#         [1, 2, 1, 2, 1, 2, 1, 2, 1, 2]])

方案2：你提到的补零相加思路实现 #

按你设想的逻辑落地代码如下：

# 初始化和目标形状一致的全零矩阵，交错后列数为原始矩阵列数的2倍
A_pad = torch.zeros(A.shape[0], 2*A.shape[1], dtype=A.dtype)
B_pad = torch.zeros(B.shape[0], 2*B.shape[1], dtype=B.dtype)

# 给A的偶数索引位赋值，B的奇数索引位赋值
A_pad[:, ::2] = A
B_pad[:, 1::2] = B

C = A_pad + B_pad

两种方案都可以得到正确的交错拼接结果，方案1代码更简洁，运算量也更低。

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