LSTM长短期记忆网络：从数学原理到PyTorch实战

1次阅读

没有评论

共计 3601 个字符，预计需要花费 10 分钟才能阅读完成。

假设我们要预测未来 3 天的股价，使用过去 30 天的收盘价作为输入。简单滑动窗口方法会面临两个核心问题：

长期依赖：第 15 天的异常波动可能影响第 28 天的走势，但传统 RNN 难以捕捉这种跨时间步的关联
梯度消失：反向传播时，误差信号需要跨越 30 个时间步传递，普通 RNN 的梯度会指数级衰减

另一个典型案例是文本生成任务。当生成第 20 个单词时，模型需要记住文章开头的主题（比如 ” 科幻小说 ”），同时遗忘无关信息（比如前文提到的配角早餐内容）。这种选择性记忆机制正是 LSTM 的设计初衷。

传统 RNN 的隐状态更新公式为：

$$h_t = \tanh(W_{xh}x_t + W_{hh}h_{t-1} + b_h)$$

其梯度可以表示为：

$$\frac{\partial h_t}{\partial h_{t-1}} = W_{hh}^T \cdot \text{diag}(1 – \tanh^2(z))$$

当时间步距 $k$ 较大时，梯度会趋近于 0（若 $W_{hh}$ 特征值 <1）或爆炸（若 $W_{hh}$ 特征值 >1）

LSTM 通过三个门控单元解决上述问题：

遗忘门（决定丢弃哪些信息）：
$$f_t = \sigma(W_f \cdot [h_{t-1}, x_t] + b_f)$$
输入门（决定存储哪些新信息）：
$$i_t = \sigma(W_i \cdot [h_{t-1}, x_t] + b_i)$$
$$\tilde{C}t = \tanh(W_C \cdot [h, x_t] + b_C)$$
输出门（决定输出什么信息）：
$$o_t = \sigma(W_o \cdot [h_{t-1}, x_t] + b_o)$$

最终细胞状态更新公式：
$$C_t = f_t \odot C_{t-1} + i_t \odot \tilde{C}_t$$
$$h_t = o_t \odot \tanh(C_t)$$

与 GRU 的显著区别在于：LSTM 维持了独立的细胞状态 $C_t$ 和隐状态 $h_t$，而 GRU 将二者合并。这使得 LSTM 在超长序列任务中表现更稳定。

class SequenceDataset(Dataset):
    def __init__(self, raw_data, seq_len=30):
        self.data = torch.FloatTensor(raw_data)
        self.seq_len = seq_len

    def __len__(self):
        return len(self.data) - self.seq_len

    def __getitem__(self, idx):
        # 输入维度: [seq_len, feature_dim]
        x = self.data[idx:idx+self.seq_len]
        # 输出维度: [predict_steps, feature_dim]
        y = self.data[idx+1:idx+self.seq_len+1]  
        return x, y

# 标准化处理
def normalize(data):
    mean = data.mean(axis=0)
    std = data.std(axis=0)
    return (data - mean) / (std + 1e-8)

class PeepholeLSTM(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size

        # 增加 peephole 连接的权重
        self.W_ci = nn.Parameter(torch.randn(hidden_size))
        self.W_cf = nn.Parameter(torch.randn(hidden_size))
        self.W_co = nn.Parameter(torch.randn(hidden_size))

        # 标准 LSTM 参数
        self.W_xi = nn.Linear(input_size, hidden_size)
        self.W_hi = nn.Linear(hidden_size, hidden_size, bias=False)
        # 类似地定义其他门参数...

    def forward(self, x, states):
        h_prev, c_prev = states

        # 带 peephole 的遗忘门
        f_t = torch.sigmoid(self.W_xf(x) + self.W_hf(h_prev) + self.W_cf * c_prev
        )
        # 其他门计算类似...

        c_t = f_t * c_prev + i_t * c_hat_t
        h_t = o_t * torch.tanh(c_t)

        return h_t, (h_t, c_t)

关键张量维度说明：

输入 x: [batch_size, seq_len, input_size]
隐状态 h: [batch_size, hidden_size]
细胞状态 c: [batch_size, hidden_size]

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for x, y in dataloader:
    optimizer.zero_grad()
    loss = model(x, y)
    loss.backward()

    # 梯度裁剪阈值通常取 5.0
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)

    optimizer.step()

假设：

批大小 batch_size=64
序列长度 seq_len=100
隐藏层维度 hidden_size=512
数据类型 float32（4 字节）

内存消耗主要来自：

正向 LSTM 的隐状态：64×100×512×4 ≈ 12.5MB
反向 LSTM 的隐状态：相同大小
中间变量缓存：约 3 倍显存

总需求 ≈ (12.5×2)×3 = 75MB 每层

使用 torch.backends.cudnn.benchmark = True 启用自动寻找最优算法
当序列长度不一致时，使用nn.utils.rnn.pack_padded_sequence

混合精度训练：

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    output = model(input)
    loss = criterion(output, target)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

采用梯度同步的 DataParallel 模式：

model = nn.DataParallel(model, device_ids=[0,1])
# 前向传播自动拆分 batch 到各 GPU
output = model(input)
# 反向传播自动聚合梯度
loss.backward()

更推荐使用 DistributedDataParallel（DDP），它在每个 GPU 上维护独立的模型副本，通过 NCCL 后端同步梯度。

测试 FP32 到 INT8 转换后的精度损失：

# 准备校准数据
calib_data = get_representative_samples()

# 量化模型
quantized_model = torch.quantization.quantize_dynamic(model, {nn.LSTM, nn.Linear}, dtype=torch.qint8
)

# 测试误差变化
fp32_err = evaluate(model, test_loader)
int8_err = evaluate(quantized_model, test_loader)
print(f"精度下降: {int8_err - fp32_err:.2%}")