PyTorch自定义peephole LSTM如何实现内置nn.LSTM的Dropout效果

你的当前实现不符合PyTorch内置nn.LSTM的Dropout逻辑，无法达到相同效果。

PyTorch官方对nn.LSTM的Dropout实现有三个明确规则：

• 仅当LSTM层数num_layers > 1时Dropout才生效，单层LSTM即使传入dropout参数，训练和推理阶段都不会施加任何Dropout
• Dropout仅施加在中间层的输出位置：除最后一层外，每一层的输出序列经过Dropout后才会作为下一层的输入，最后一层的输出不会附加任何Dropout
• 使用的是变分Dropout（Variational Dropout）：对同一层内同一个batch的样本，所有时间步共享同一个Dropout掩码，不会为每个时间步单独生成随机掩码

• 无论LSTM层数多少，直接对最终返回的隐藏序列加Dropout，和官方“单层不生效、仅中间层加Dropout”的逻辑冲突
• 直接调用nn.Dropout处理形状为(seq_len, batch_size, hidden_size)的序列时，会对每个时间步、每个隐藏维度独立采样掩码，和官方“同层同样本所有时间步共享掩码”的变分Dropout逻辑不一致

1. 单层LSTM不需要添加任何Dropout逻辑，直接返回输出即可
2. 多层LSTM时，遍历每一层做前向传播，仅对非最后一层的输出施加变分Dropout，处理后再传入下一层
3. 变分Dropout不需要额外初始化nn.Dropout层，前向传播时动态生成和单步隐藏状态形状一致的掩码，广播应用到该层所有时间步的输出上即可

参考实现代码：

import torch
import torch.nn.functional as F

# 以多层peephole LSTM的初始化、前向逻辑为例
def __init__(self, input_size, hidden_size, num_layers=1, dropout=0.0, ...):
    super().__init__()
    self.num_layers = num_layers
    self.dropout_p = dropout
    # 用ModuleList存储每一层的peephole LSTM实现
    self.lstm_layers = nn.ModuleList([
        # 第一层输入维度为input_size，后续层输入维度为hidden_size
        PeepholeLSTMLayer(input_size if i==0 else hidden_size, hidden_size) 
        for i in range(num_layers)
    ])

def forward(self, x, init_states=None):
    # x默认形状为(seq_len, batch_size, input_size)
    current_input = x
    layer_h_list, layer_c_list = [], []
    
    for layer_idx, layer in enumerate(self.lstm_layers):
        # 读取当前层初始状态
        if init_states is not None:
            h_init, c_init = init_states[0][layer_idx], init_states[1][layer_idx]
            layer_hidden_seq, (h_t, c_t) = layer(current_input, (h_init, c_init))
        else:
            layer_hidden_seq, (h_t, c_t) = layer(current_input)
        
        layer_h_list.append(h_t)
        layer_c_list.append(c_t)
        
        # 仅非最后一层、训练模式、dropout概率大于0时施加变分Dropout
        if layer_idx != self.num_layers - 1 and self.training and self.dropout_p > 0:
            # 生成和单步隐藏状态形状一致的掩码：(batch_size, hidden_size)
            dropout_mask = F.dropout(
                torch.ones_like(h_t, device=x.device, dtype=x.dtype),
                p=self.dropout_p,
                training=True
            )
            # 广播掩码到所有时间步：掩码形状变为(1, batch_size, hidden_size)，自动对齐序列长度维度
            layer_hidden_seq = layer_hidden_seq * dropout_mask.unsqueeze(0)
        
        current_input = layer_hidden_seq
    
    final_hidden_seq = current_input
    final_h = torch.stack(layer_h_list, dim=0)
    final_c = torch.stack(layer_c_list, dim=0)
    return final_hidden_seq, (final_h, final_c)

不需要为每一层单独初始化nn.Dropout实例，上述动态生成掩码的方式已经对齐官方实现的缩放逻辑：保留的权重会自动除以1-p，保证输出期望和推理阶段一致，不需要手动做数值缩放。

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