PyTorch中model.eval()仅禁用外层类Dropout的问题及解决

问题现象 #

复现机器翻译Transformer模型时，测试阶段已调用model.eval()和torch.no_grad()执行推理，但输入相同序列时输出结果不一致。通过检查Dropout层的training属性：

print(model.dropout_enc.training)
print(model.encoder_layers[0].dropout_multi_head.training)

得到输出：

False
True

可见外层Dropout已被禁用，但子层内的Dropout仍处于训练状态。

代码背景 #

Transformer类初始化时用普通列表创建子层：

class Transformer(nn.Module):
    def __init__(self, vocab_size_in, vocab_size_out, embedding_dim, n_heads, key_dim, value_dim, ffn_dim, n=10000,
                 eps=1e-5, padding_token_index=0, p_drop=0.1, n_encoder_layers=1, n_decoder_layers=1):
        super(Transformer, self).__init__()
        # ...
        self.dropout_enc = nn.Dropout(p_drop)
        self.encoder_layers = [EncoderLayer(embedding_dim, key_dim, value_dim, ffn_dim, n_heads, p_drop, eps)] * n_encoder_layers
        # ...
        self.dropout_dec = nn.Dropout(p_drop)
        self.decoder_layers = [DecoderLayer(embedding_dim, key_dim, value_dim, ffn_dim, n_heads, p_drop, eps)] * n_decoder_layers
        # ...

EncoderLayer内部包含Dropout层：

class EncoderLayer(nn.Module):
    def __init__(self, embedding_dim=512, key_dim=512, value_dim=512, ffn_dim=512, n_heads=8, p_drop=0.1, eps=1e-5):
        super().__init__()
        # ...
        self.dropout_multi_head = nn.Dropout(p_drop)
        # ...
        self.dropout_FFN = nn.Dropout(p_drop)
        # ...

原因分析 #

使用普通列表[Module] * num_layers创建子层时，这些子模块没有被PyTorch的nn.Module系统管理，model.eval()只会递归处理直接注册为模块属性的子模块（如self.dropout_enc），不会遍历普通列表里的子层，导致子层内的Dropout仍保持训练模式，从而在推理时产生随机结果。

解决方案 #

用nn.ModuleList替代普通列表初始化子层，PyTorch会自动将列表内的所有子模块纳入管理，调用model.eval()时会递归切换所有子模块到评估模式：

class Transformer(nn.Module):
    def __init__(self, vocab_size_in, vocab_size_out, embedding_dim, n_heads, key_dim, value_dim, ffn_dim, n=10000,
                 eps=1e-5, padding_token_index=0, p_drop=0.1, n_encoder_layers=1, n_decoder_layers=1):
        super(Transformer, self).__init__()
        # ...
        self.dropout_enc = nn.Dropout(p_drop)
        # 改用nn.ModuleList创建编码器层
        self.encoder_layers = nn.ModuleList([
            EncoderLayer(embedding_dim, key_dim, value_dim, ffn_dim, n_heads, p_drop, eps) 
            for _ in range(n_encoder_layers)
        ])
        # ...
        self.dropout_dec = nn.Dropout(p_drop)
        # 改用nn.ModuleList创建解码器层
        self.decoder_layers = nn.ModuleList([
            DecoderLayer(embedding_dim, key_dim, value_dim, ffn_dim, n_heads, p_drop, eps) 
            for _ in range(n_decoder_layers)
        ])
        # ...

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