LSTM模型定义（PyTorch）

• LSTM（Long Short-Term Memory）长短期记忆网络
  是 RNN（循环神经网络）的一种改进版本，主要用来解决 时间序列预测 或 需要记住过去信息的任务，例如：👇🏻
  • PM2.5 时间序列预测
  • 文本生成
  • 股票预测
  • 温度预测
  • 电力负载预测
• 普通 RNN 的问题是运行久了就遗忘前面的信息（梯度消失），而 LSTM 通过 “门结构（gates）” 让网络能够选择：👇🏻
  • 记住（Keep）
  • 忘掉（Forget）
  • 更新（Update）
• 这些信息。

import torch
import torch.nn as nn


# 定义一个用于回归任务的 LSTM 模型
class LSTMRegressor(nn.Module):
    def __init__(self, input_dim, hidden_dim=64, num_layers=2, dropout=0.2, bidirectional=False):
        super().__init__()

        # 保存模型参数（不是必须，但方便调试或扩展示例）
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.bidirectional = bidirectional

        # 定义 LSTM 网络层
        # input_size = 输入的特征维度（例如每小时1个 PM2.5，则=1）
        # hidden_size = LSTM 隐藏层维度
        # num_layers = 堆叠的 LSTM 层数（深度）
        # batch_first = True，让输入输出维度变成 (batch, seq_len, feature)
        # dropout = 每层之间的 dropout（LSTM 内部不使用 dropout）
        # bidirectional = 是否使用双向 LSTM（会将 forward 和 backward 拼接）
        self.lstm = nn.LSTM(input_size=input_dim,
                            hidden_size=hidden_dim,
                            num_layers=num_layers,
                            batch_first=True,
                            dropout=dropout if num_layers > 1 else 0.0,
                            bidirectional=bidirectional)

        # 如果是双向 LSTM，则输出特征维度需要 ×2
        direction = 2 if bidirectional else 1

        # 定义全连接回归头（用于将 LSTM 输出映射到标量）
        # LSTM 输出维度 = hidden_dim * direction
        # 先压缩到 hidden_dim//2，再 ReLU 激活，最后输出 1 个值
        self.fc = nn.Sequential(
            nn.Linear(hidden_dim * direction, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, 1)
        )

    def forward(self, x):
        """
        x 的形状：(batch_size, seq_len, input_dim)
        """
        # out：每个时间步的输出 -> 维度 (batch, seq_len, hidden_dim*direction)
        # (h_n, c_n)：最后一层每个方向的最后隐藏态、细胞态（不用也可以）
        out, (h_n, c_n) = self.lstm(x)

        # 取最后一个时间步的输出作为序列特征
        # last: (batch_size, hidden_dim * direction)
        last = out[:, -1, :]

        # 输入全连接网络，得到最后的回归预测（batch_size, 1）
        return self.fc(last)