PyTorch LSTM训练时输出与目标张量维度不匹配问题求助

问题根源 #

你的LSTM当前输出是torch.Size([64,20,64])，这是因为batch_first=True时，LSTM会返回整个序列每个时间步的隐藏状态（维度为[批量大小, 序列长度, 隐藏层维度]），但你的任务是单步预测——用20步窗口数据预测窗口后的第21步股价，目标形状是[64,1]，两者维度完全不匹配，硬扩展目标只会导致错误的损失计算逻辑。

解决步骤 #

1. 重构LSTM模型，添加全连接层映射到目标维度 #

需要从LSTM的序列输出中提取最后一个时间步的结果，再通过全连接层把隐藏层的64维特征映射到1维预测值：

import torch.nn as nn
import torch.optim as optim

class LSTMStockPredictor(nn.Module):
    def __init__(self, input_size=1, hidden_size=64, num_layers=2):
        super().__init__()
        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, 1)  # 将64维隐藏状态转为1维预测值

    def forward(self, x):
        # x输入形状: [batch_size, seq_len, input_size] = [64, 20, 1]
        lstm_out, _ = self.lstm(x)  # lstm_out形状: [64, 20, 64]
        # 提取最后一个时间步的输出，对应窗口结束后的预测节点
        last_step_out = lstm_out[:, -1, :]  # 形状变为[64, 64]
        output = self.fc(last_step_out)  # 最终输出形状: [64, 1]
        return output

2. 替换原模型初始化代码 #

model = LSTMStockPredictor()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

3. 验证训练时的形状匹配 #

训练循环中，output的形状会是[64,1]，和目标y的[64,1]完全一致，不会再出现形状不匹配的警告。

额外检查点 #

• 确认数据集构造逻辑：确保stock_train_tensor[i]对应的是连续20步的股价数据，stock_train_target_tensor[i]确实是这20步之后的第1步股价，数据对应关系没有错误。
• 如果你的任务是多步预测（比如用20步预测接下来5步），则需要调整目标张量的形状为[批量大小, 预测步数]，并修改模型输出为对应步数的结果，但根据你的目标张量形状，当前任务是单步预测，按上述方案即可解决。

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