PyTorch中如何拼接形状不同的两模型输出以输入LSTM？

你的核心问题是ModelB带有空间维度（height, width），无法直接和ModelA的输出拼接，需要先把空间信息压缩/转换为特征维度的一部分，再对齐两个模型的维度顺序，最后拼接输入LSTM。以下是具体步骤：

1. 处理ModelB的空间维度 #

ModelB输出的[height, width]是空间特征图，必须把这部分信息整合到特征维度里，常用三种方法：

方法一：全局池化（最常用，轻量高效） #

用全局平均/最大池化把每个空间特征图压缩成单个值，直接对齐ModelA的结构：

import torch
import torch.nn as nn

# 假设model_b_out是ModelB的输出，形状[batch, feat_b, frames, h, w]
# 先调整维度顺序，把frames放到第2位（方便后续处理）
model_b_reshaped = model_b_out.permute(0, 2, 1, 3, 4)  # [batch, frames, feat_b, h, w]
# 全局平均池化，把h和w压缩成1
gap = nn.AdaptiveAvgPool2d(1)
model_b_processed = gap(model_b_reshaped)
# 去掉最后两个多余的维度，得到[batch, frames, feat_b]
model_b_processed = model_b_processed.squeeze(-1).squeeze(-1)

方法二：展平空间维度（保留全部空间信息，适合小尺寸特征图） #

直接把h*w的空间信息展平成特征的一部分，缺点是如果h/w较大，会导致特征维度暴增：

# 同样先调整维度顺序
model_b_reshaped = model_b_out.permute(0, 2, 1, 3, 4)  # [batch, frames, feat_b, h, w]
# 展平feat_b、h、w为一个特征维度
model_b_processed = model_b_reshaped.flatten(2)  # [batch, frames, feat_b*h*w]

方法三：卷积降维（平衡信息保留和维度大小） #

用1x1卷积先压缩空间特征的通道数，再展平，比池化保留更多空间细节，比直接展平更可控：

# 先合并batch和frames维度，方便卷积处理
model_b_reshaped = model_b_out.permute(0, 2, 1, 3, 4).flatten(0, 1)  # [batch*frames, feat_b, h, w]
# 用1x1卷积把通道数调整到目标值（比如256）
conv = nn.Conv2d(model_b_out.size(1), 256, kernel_size=1)
model_b_conv = conv(model_b_reshaped)
# 展平特征和空间维度
model_b_flattened = model_b_conv.flatten(1)  # [batch*frames, 256*h*w]
# 恢复batch和frames维度
model_b_processed = model_b_flattened.unflatten(0, (model_b_out.size(0), model_b_out.size(2)))  # [batch, frames, 256*h*w]

2. 调整ModelA的维度顺序 #

ModelA的输出是[batch, feat_a, frames]，需要把frames放到序列维度（第2位），和处理后的ModelB对齐：

# model_a_out是ModelA的输出，形状[batch, feat_a, frames]
model_a_processed = model_a_out.permute(0, 2, 1)  # [batch, frames, feat_a]

3. 拼接特征并输入LSTM #

现在两个输出都是[batch, frames, feature_dim]的形状，直接在最后一维拼接，然后输入LSTM（记得设置batch_first=True匹配输入形状）：

# 拼接特征，得到[batch, frames, feat_a + feat_b(或处理后的特征数)]
concat_feat = torch.cat([model_a_processed, model_b_processed], dim=-1)

# 初始化LSTM，input_size是拼接后的特征维度
lstm = nn.LSTM(input_size=concat_feat.size(-1), hidden_size=256, batch_first=True)
# 输入LSTM，得到输出和隐藏状态
lstm_out, (h_n, c_n) = lstm(concat_feat)

关键注意点 #

• 特征维度平衡：如果两个模型处理后的特征维度差距过大（比如一个256，一个上万），会导致LSTM偏向学习维度大的特征，建议用全连接层或卷积把其中一个的特征维度调整到相近范围。
• 帧数量对齐：必须保证两个模型输出的num of frames完全一致，如果不一致，要先做帧裁剪或插值对齐。
• 预训练模型冻结：如果不需要微调ModelA/ModelB，记得设置for param in model_a.parameters(): param.requires_grad = False，避免反向传播更新预训练权重。

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