PyTorch LSTM训练异常：无法利用历史数据，损失停滞于0.25

我使用PyTorch的LSTM API构建测试模型，任务规则为：当前数值大于前一个时输出1，否则输出0，首个输出固定为1.0。例如输入数组[0.7, 0.3, 0.9, 0.99]，预期输出为[1.0, 0.0, 1.0, 1.0]。但训练后损失始终维持在0.25左右不再下降，模型似乎一直在输出0和1的平均值0.5，无法利用历史数据完成任务。

网络结构代码 #

# network.py

import torch

N_INPUT = 1
N_STACKS = 1
N_HIDDEN = 3

LR = 0.001


class Network(torch.nn.Module):

    def __init__(self):
        super(Network, self).__init__()

        self.lstm = torch.nn.LSTM(
            input_size=N_INPUT,
            hidden_size=N_HIDDEN,
            num_layers=N_STACKS,
        )

        self.linear = torch.nn.Linear(N_HIDDEN, 1)
        self.relu = torch.nn.ReLU()

        self.optim = torch.optim.Adam(self.parameters(), lr=LR)
        self.loss = torch.nn.MSELoss()

    def backprop(self, xs, es):
        self.optim.zero_grad()
        l = self.loss(xs, torch.tensor(es))
        l.backward()
        self.optim.step()
        return l

    def forward(self, dat):
        out, _ = self.lstm(torch.tensor(dat))
        out = self.relu(out)
        out = self.linear(out)
        return out

训练代码 #

# main.py

import network
import numpy as np

n: network.Network = network.Network()

def rand_array():
    a = [[np.random.uniform(0, 1)] for i in range(1000)]
    expected = [0.0 if a[i - 1][0] > a[i][0] else 1.0 for i in range(len(a))]
    expected[0] = 1.0
    return [a, expected]

data = [rand_array() for i in range(1000)]

for i in range(100):
    for i in data:
        pred = n(i[0])
        loss = n.backprop(pred, i[1])
        print("Loss: {:.5f}".format(loss))

问题原因分析 #

• LSTM输入序列格式错误：PyTorch的LSTM默认期望输入形状为(seq_len, batch_size, input_size)，但当前代码中传入的torch.tensor(dat)形状为(1000, 1)，会被LSTM解析为batch_size=1000, seq_len=1。这意味着每个数据点被当作独立的单步序列处理，LSTM无法捕捉序列的连续性，自然无法记住前一个时间步的数值。
• 输出层缺少合适的激活函数：任务是0/1二分类，当前输出层仅用线性层，输出范围无限制，模型容易收敛到0和1的平均值0.5（此时MSE损失恰好为0.25），无法输出符合任务要求的离散值。
• 张量形状不匹配：模型输出pred的形状为(1000, 1)，而预期输出es转成张量后形状为(1000,)，虽然PyTorch会自动广播，但这种不匹配可能导致损失计算逻辑异常，影响梯度传递。
• 训练循环变量名冲突：外层循环使用i作为epoch变量，内层循环又用i遍历数据，会导致变量覆盖，可能打乱训练流程。

修复方案 #

1. 调整LSTM输入格式，启用序列处理 #

在Network.forward中，将输入数据调整为LSTM要求的序列格式，同时启用batch_first=True简化维度管理：

# 修改Network.__init__中的LSTM定义
self.lstm = torch.nn.LSTM(
    input_size=N_INPUT,
    hidden_size=N_HIDDEN,
    num_layers=N_STACKS,
    batch_first=True  # 启用batch_first，输入形状变为(batch_size, seq_len, input_size)
)

# 修改forward方法
def forward(self, dat):
    # 将输入从(seq_len, input_size)转为(1, seq_len, input_size)，代表1个batch，长度为seq_len的序列
    x = torch.tensor(dat, dtype=torch.float32).unsqueeze(0)
    out, _ = self.lstm(x)
    out = self.relu(out)
    out = self.linear(out)
    out = torch.sigmoid(out)  # 添加sigmoid激活，将输出限制在0-1区间
    return out.squeeze(0)  # 移除batch维度，返回(seq_len, 1)的结果

2. 修正损失计算的张量形状和类型 #

在backprop方法中，确保预期输出的形状和数据类型与模型输出一致：

def backprop(self, xs, es):
    self.optim.zero_grad()
    # 将预期输出转为float32张量，并调整为与xs相同的形状(seq_len, 1)
    target = torch.tensor(es, dtype=torch.float32).unsqueeze(1)
    l = self.loss(xs, target)
    l.backward()
    self.optim.step()
    return l

3. 修复训练循环的变量名冲突 #

# 修改训练循环
for epoch in range(100):
    total_loss = 0.0
    for item in data:
        pred = n(item[0])
        loss = n.backprop(pred, item[1])
        total_loss += loss.item()
    # 每轮epoch打印平均损失，避免频繁输出
    print(f"Epoch {epoch+1}, Average Loss: {total_loss/len(data):.5f}")

4. 可选：调整学习率 #

当前学习率0.001可能偏小，可尝试提升至0.01，加速模型收敛：

LR = 0.01

验证效果 #

修改后，模型将正确处理序列数据，利用LSTM的记忆能力对比当前值与前一个值，sigmoid激活会将输出限制在0-1区间，损失会逐渐下降至接近0的水平，模型输出也会符合任务预期。

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