LSTM模型输出异常：批量输入仅得单个输出，需每行对应输出

训练集维度为[7000, 2]，但当前LSTM模型仅输出单个数值，期望配置模型后，输入7000行2列的数据时，输出7000行1列的结果，即每一行输入对应一个输出。

问题根源 #

1. 模型forward方法中仅取了LSTM输出的最后一个时间步结果：out = self.fc(out[:, -1, :])，导致仅输出单个预测值，而非每个时间步都生成输出。
2. 输入数据的形状处理不当，手动设置的初始隐藏状态与实际batch size不匹配，进一步导致输出维度不符。

解决方案 #

1. 修改模型的forward逻辑 #

将全连接层应用到LSTM的所有时间步输出上，而非仅最后一个时间步：

def forward(self, x, h0=None, c0=None):
    if h0 is None or c0 is None:
        h0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim).to(x.device)
        c0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim).to(x.device)
    
    out, (hn, cn) = self.lstm(x, (h0, c0))
    # 对所有时间步的输出应用全连接层
    out = self.fc(out)
    return out, hn, cn

2. 调整输入数据形状 #

将输入从[7000,2]转换为[7000,1,2]，让每一行作为独立的单步时间序列输入模型：

outputs, h0, c0 = model(X_train.unsqueeze(1), h0, c0)

3. 修正隐藏状态初始化 #

去掉手动设置的固定形状隐藏状态，让模型根据输入的batch size自动初始化，避免维度不匹配：

h0, c0 = None, None  # 由模型自动生成匹配batch size的隐藏状态

完整修改后的代码 #

import torch
import torch.nn.functional as F
import torch.nn as nn

class LSTMModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, layer_dim, output_dim,  batch_first=True):
        super(LSTMModel, self).__init__()
        self.hidden_dim = hidden_dim
        self.layer_dim = layer_dim
        self.lstm = nn.LSTM(input_dim, hidden_dim, layer_dim, batch_first=batch_first)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x, h0=None, c0=None):
        if h0 is None or c0 is None:
            h0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim).to(x.device)
            c0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim).to(x.device)
        
        out, (hn, cn) = self.lstm(x, (h0, c0))
        # 应用全连接层到所有时间步输出
        out = self.fc(out)
        return out, hn, cn

# 补充缺失的全局变量定义
ROWS = 7000
FEATS = 2
TRAIN_FRAC = 1.0  # 可根据实际训练比例调整

g = torch.Generator().manual_seed(2147483647) # 保证实验可复现
X  = torch.randn((ROWS, FEATS), generator=g)
Y = X[:,0] + X[:,1].roll(1)
Y[0] = X[0,0]
print(X[:10], '\n------------\n', Y[:10], '\n', X.shape, Y.shape)

last_train_ix = int(TRAIN_FRAC*ROWS)

X_train = X[:last_train_ix,]
Y_train = Y[:last_train_ix,]


model = LSTMModel(input_dim=2, hidden_dim=100, layer_dim=1, output_dim=1, batch_first=True)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

num_epochs = 100
h0, c0 = None, None  # 自动初始化隐藏状态

for epoch in range(num_epochs):
    model.train()
    optimizer.zero_grad()
    # 调整输入形状为[7000,1,2]
    outputs, h0, c0 = model(X_train.unsqueeze(1), h0, c0)
    # 压缩维度匹配Y_train的形状
    outputs = outputs.squeeze()
    print(X_train.shape, outputs.shape, Y_train.shape)
    
    loss = criterion(outputs, Y_train)
    loss.backward()
    optimizer.step()

    h0 = h0.detach()
    c0 = c0.detach()

    if (epoch+1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

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