如何在PyTorch中拼接不同形状的两种模态输入张量

import torch
a = torch.randn((1, 30, 1220)) # 文本嵌入：1个批次，30个span，每个维度1220
b = torch.randn((1, 128, 256)) # 图像特征：1个批次，128个特征块，每个维度256

问题1：将b的全部内容拼接到a的30个span中的每一个上 #

实现逻辑为：先把b的所有特征展平为一维向量，再复制30份匹配a的span数量，最后在特征维度完成拼接。

# 1. 展平b的所有特征，形状从(1,128,256)转为(1, 32768)
b_flatten = b.flatten(start_dim=1)
# 2. 扩展出span维度，重复b的内容匹配30个span，形状转为(1,30,32768)
b_expand = b_flatten.unsqueeze(1).repeat(1, 30, 1)
# 3. 特征维度拼接
result1 = torch.cat([a, b_expand], dim=-1)
print(result1.shape) # 输出：torch.Size([1, 30, 33988])

问题2：将完整的b和完整的a直接拼接 #

根据下游任务需求，有两种常用实现方案：

场景1：按序列维度拼接（适配序列建模任务，比如后续接Transformer） #

需要先把图像特征维度投影到和文本特征一致，再在序列长度维度拼接：

# 线性投影层对齐特征维度
proj_layer = torch.nn.Linear(256, 1220)
b_proj = proj_layer(b)
# 序列维度拼接
result_seq = torch.cat([a, b_proj], dim=1)
print(result_seq.shape) # 输出：torch.Size([1, 158, 1220])

场景2：全局特征拼接（适配分类/回归等任务） #

直接将两个张量全部展平后拼接：

a_flat = a.flatten(start_dim=1)
b_flat = b.flatten(start_dim=1)
result_global = torch.cat([a_flat, b_flat], dim=-1)
print(result_global.shape) # 输出：torch.Size([1, 69368])

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