如何在PyTorch中实现LSTM网络遗忘门偏置初始化以提升模型性能

首先明确PyTorch内置LSTM的偏置存储规则：

• 每个LSTM层包含两组偏置参数：bias_ih（输入特征到隐藏层的偏置）、bias_hh（上一步隐藏状态到当前隐藏层的偏置）
• 两组偏置的长度均为4 * 隐藏层维度，对应四个门控的偏置按顺序拼接：输入门、遗忘门、细胞更新门、输出门
• 遗忘门的偏置对应区间为[hidden_size : 2 * hidden_size]

具体实现代码 #

要将遗忘门偏置初始化为常用的1（避免训练初期遗忘过多历史信息），可参考以下代码：

import torch
import torch.nn as nn

# 示例：定义LSTM模型，输入维度128，隐藏层维度64，层数为2
lstm_layer = nn.LSTM(input_size=128, hidden_size=64, num_layers=2, batch_first=True)

# 遍历所有参数修改遗忘门偏置
for param_name, param in lstm_layer.named_parameters():
    if "bias" in param_name:
        single_gate_size = param.shape[0] // 4
        # 给遗忘门对应位置的偏置赋值为1
        nn.init.constant_(param[single_gate_size : 2*single_gate_size], 1.0)

如果需要自定义其他初始化值，直接修改nn.init.constant_的第二个参数即可，也可以替换为其他初始化函数适配不同需求。

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