PyTorch下Seq2Seq LSTM模型无法正常训练的问题排查

问题背景 #

我正在用PyTorch实现Seq2Seq LSTM模型解决序列预测任务：输入长度为5的序列，预测后续长度为5的序列。数据张量尺寸为[bs, seq_length, features]（seq_length=10，features=1），每个特征是0~3的整数。目前模型无法正常训练，怀疑问题出在数据转换环节，尝试过MinMaxScaler归一化到[0,1]、在Linear层后加sigmoid，但5000轮训练后曲线仍异常（波动极大，无明显下降趋势）。

模型定义代码 #

class LSTM(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers, dropout_prob):
        super(LSTM, self).__init__()
        self.lstm_layer = nn.LSTM(input_dim, hidden_dim, num_layers, dropout=dropout_prob)
        self.output_layer = nn.Linear(hidden_dim, output_dim)

    def forward(self, X):
        out, (hidden, cell) = self.lstm_layer(X)
        out = self.output_layer(out)
        return out

训练循环代码 #

def train_loop(t, checkpoint_epoch, dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, X in enumerate(dataloader):
        X = X[0].type(torch.float).to(device)

        # X = torch.Size([batch_size, 10, input_dim])
        # Split sequences into input and target
        inputs = transform(X[:, :5, :]) # inputs = [batch_size, 5, input_dim]
        targets = transform(X[:, 5:, :]) # targets = [batch_size, 5, input_dim]

        # predictions (forward pass)
        with autocast():
            pred = model(inputs)  # pred = [batch_size, 5, input_dim]
            loss = loss_fn(pred, targets)

        # backprop
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            #print(f"Current loss: {loss:>7f}, [{current:>5d}/{size:>5d}]")

        # Delete variables and empty cache
        del X, inputs, targets, pred
        torch.cuda.empty_cache()

    return loss

数据预处理代码 #

def main():
    num_agents = 2
    # Open the HDF5 file
    with h5py.File('dataset_' + str(num_agents) + 'UAV.hdf5', 'r') as f:
        # Access the dataset
        data = f['data'][:]
        # Convert to PyTorch tensor
        data_tensor = torch.tensor(data)

        size = data_tensor.size()
        seq_length = 10
        reshaped = data_tensor.view(-1, size[2], size[3])

        r_size = reshaped.size()
        reshaped = reshaped[:, :, 1:]
        reshaped_v2 = reshaped.view(r_size[0], -1)

        dataset = create_dataset(reshaped_v2.numpy(), seq_length)

        f.close()

    dataset = TensorDataset(dataset)

    # Split the dataset into training and validation sets
    train_size = int(0.8 * len(dataset))  # 80% for training
    val_size = len(dataset) - train_size  # 20% for validation
    train_dataset, val_dataset = random_split(dataset, [train_size, val_size])

    train_dataloader = DataLoader(train_dataset, batch_size=params['batch_size'], shuffle=True, pin_memory=True)
    val_dataloader = DataLoader(val_dataset, batch_size=params['batch_size'], shuffle=False, pin_memory=True)

1. 训练流程中的问题 #

• 数据预处理逻辑混乱：
  预处理中的多次view和维度裁剪（reshaped[:, :, 1:]）可能破坏了原始数据的时序连续性，比如原本的多智能体时序数据被强行展平，导致生成的训练样本是无意义的序列，模型根本无法学习到有效模式。
• 归一化使用错误：
  对输入和目标都应用transform（MinMaxScaler）时，若scaler是基于整个数据集拟合的，会导致数据泄露，模型无法泛化；另外，目标是离散的0~3整数，归一化后变成连续浮点数，和模型输出的匹配度差。
• LSTM输入维度顺序错误：
  PyTorch的nn.LSTM默认输入维度为[seq_len, batch_size, input_dim]，但你的输入inputs是[batch_size, 5, input_dim]，这会让LSTM错误地将batch维度当成序列长度处理，完全偏离预期。
• 冗余的内存操作：
  训练循环中的del变量和torch.cuda.empty_cache()属于多余操作，PyTorch会自动管理内存，频繁清空缓存反而会导致显存波动、训练变慢。
• 缺乏验证监控：
  只跟踪训练损失，无法判断模型是未学习还是过拟合，没有验证集损失的对比，无法及时调整策略。

2. 认知错误 #

• 对LSTM输出尺度的误解：
  LSTM隐藏层输出是tanh激活的[-1,1]范围，但经过Linear层后，输出没有激活限制，尺度是任意的，并非固定在[-1,1]。你之前认为pred是[-1,1]是错误的，不需要强行加sigmoid适配，除非任务明确要求输出[0,1]。
• 离散序列用回归损失的误区：
  你的数据是0~3的离散整数，属于序列分类任务，而非回归任务。用MSE等回归损失会让模型学习连续值，无法适配离散的类别输出，应该用交叉熵损失。
• 对Seq2Seq结构的误解：
  当前模型只是普通LSTM+Linear的直接映射，并非真正的Seq2Seq模型。标准Seq2Seq包含编码器（处理输入序列）和解码器（逐步生成预测序列），若需要自回归生成（用前一个预测结果生成下一个），当前结构无法满足。
• 归一化的适用场景误解：
  离散整数序列不需要归一化，归一化会把离散的类别变成连续浮点数，反而干扰模型学习离散模式。如果是分类任务，直接用原始整数作为标签即可。

修正建议 #

1. 修复数据预处理：
   确保create_dataset生成的样本是时序连续的，每个样本为长度10的序列，前5个为输入、后5个为目标，避免无意义的维度转换。
2. 调整模型结构：
   • 给LSTM添加batch_first=True，匹配你的输入维度：

     self.lstm_layer = nn.LSTM(input_dim, hidden_dim, num_layers, dropout=dropout_prob, batch_first=True)
     

   • 若为分类任务，将输出层改为对应4个类别的线性层，用交叉熵损失：

     self.output_layer = nn.Linear(hidden_dim, 4)
     # 损失函数使用nn.CrossEntropyLoss()，target需转为torch.LongTensor类型
     

3. 修正数据处理逻辑：
   • 分类任务无需归一化，直接用原始整数作为target；若坚持回归任务，scaler仅拟合训练集，验证/测试集用训练集的scaler转换。
4. 优化训练循环：
   • 删除冗余的del和torch.cuda.empty_cache()操作；
   • 加入验证集损失计算，每轮训练后评估验证集，监控训练/验证损失变化。
5. 调整超参数：
   • 初始学习率设为1e-4，用Adam优化器；
   • 隐藏层维度选64/128，层数1~2层，dropout不超过0.3。

内容的提问来源于stack exchange，提问作者Enrique Vilchez Campillejo