PyTorch中如何无循环实现两个张量所有批次配对的拼接操作

你要的功能可以通过张量扩维+广播拼接一行实现，完全不需要嵌套循环，性能远高于Python层循环写法，适配任意批次规模：

import torch

# 你原来的张量定义
S = torch.rand((3,2,1))
T = torch.ones((3,2,1))

# 高效实现代码
C = torch.cat([S.unsqueeze(1), T.unsqueeze(0)], dim=2)

# 验证形状符合要求：输出torch.Size([3, 3, 4, 1])
print(C.shape)

实现原理 #

• S.unsqueeze(1)：给S新增第1维，形状从(3,2,1)变为(3,1,2,1)
• T.unsqueeze(0)：给T新增第0维，形状从(3,2,1)变为(1,3,2,1)
• 两个张量会自动在新增的维度上广播，都扩展为(3,3,2,1)的形状
• 沿第2维拼接后，第2维长度变为2+2=4，最终得到形状为(3,3,4,1)的结果，和你嵌套循环的计算结果完全一致，你可以通过torch.allclose(C, 循环生成的C)验证正确性。

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