如何在PyTorch中实现张量的指定规则交替拼接？

要实现你需要的交替拼接逻辑，可以通过张量重塑、维度扩展、拼接和展平这几步完成，具体代码如下：

import torch

# 初始化输入张量
a = torch.tensor([[[1,2,3,4,5,6,7,8,9]],[[11,12,13,14,15,16,17,18,19]]])
b = torch.tensor([[[10, 20, 30]], [[1, 2, 3]]])

# 1. 将a按第三维度每3个元素分组重塑
a_grouped = a.view(2, 1, 3, 3)
# 2. 给b添加一个维度，匹配分组后的a的维度结构
b_expanded = b.unsqueeze(-1)
# 3. 在最后一维拼接分组后的a和扩展后的b
combined = torch.cat([a_grouped, b_expanded], dim=-1)
# 4. 展平第三、第四维度得到最终结果
result = combined.flatten(start_dim=2)

print(result)

输出结果：

tensor([[[ 1,  2,  3, 10,  4,  5,  6, 20,  7,  8,  9, 30]],
        [[11, 12, 13,  1, 14, 15, 16,  2, 17, 18, 19,  3]]])

操作细节解释 #

• view(2, 1, 3, 3)：把a的第三维度（长度9）拆分为3组，每组3个元素，张量形状变为[2,1,3,3]，这样能和b的元素一一对应。
• unsqueeze(-1)：在b的最后添加一个维度，使其形状变为[2,1,3,1]，确保和分组后的a在拼接维度上兼容。
• torch.cat(..., dim=-1)：在最后一个维度拼接，让b的每个元素插入到a对应组的末尾。
• flatten(start_dim=2)：从第三维度开始展平，把[3,4]的结构合并为长度12的一维，还原成目标的张量形状。

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