在PyTorch中为预训练ResNet模型添加Dropout遇问题求助

问题描述 #

尝试为预训练ResNet的所有Conv2d/Conv1d层后添加Dropout，但运行代码后仅在模型最后出现一个Dropout，相关代码及输出如下：

原代码片段 #

feats_list = []
for key, value in model._modules.items():
    feats_list.append(value)

for feat in feats_list:
    if isinstance(feat, nn.Conv2d) or isinstance(feat, nn.Conv1d):
        feats_list.append(nn.Dropout(p=0.5, inplace=True))

model.features = nn.Sequential(*feats_list)
print(model.features)

实际输出 #

...
  (8): AdaptiveAvgPool2d(output_size=(1, 1))
  (9): Linear(in_features=2048, out_features=1000, bias=True)
  (10): Dropout(p=0.5, inplace=True)

问题原因 #

1. 未遍历子模块：ResNet的核心Conv2d层大多嵌套在layer1-layer4的BasicBlock/Bottleneck子模块中，而model._modules.items()仅返回顶层模块（如conv1、avgpool等），原代码无法检测到深层的Conv层。
2. Dropout位置错误：即使检测到顶层Conv层，代码是将Dropout追加到列表末尾，而非紧跟在对应Conv层之后，最终所有Dropout都堆积在模型最后。

解决方案 #

通过递归遍历模型的所有子模块，在每个Conv2d层后插入Dropout层。以下是修改后的代码：

递归处理模块的工具函数 #

def insert_dropout_after_conv(module, dropout_p=0.5):
    """递归遍历模块，在每个Conv2d后添加Dropout"""
    new_layers = []
    for layer in module.children():
        new_layers.append(layer)
        # 对Conv2d层添加Dropout
        if isinstance(layer, nn.Conv2d):
            new_layers.append(nn.Dropout(p=dropout_p, inplace=True))
        # 递归处理嵌套的子模块
        elif len(list(layer.children())) > 0:
            new_layers.append(insert_dropout_after_conv(layer, dropout_p))
    return nn.Sequential(*new_layers)

修改后的generic_classifier函数片段 #

def generic_classifier(model, criterion, optimizer, num_epochs=25):
    # 递归在所有Conv2d后添加Dropout
    model = insert_dropout_after_conv(model)
    
    # 替换全连接层（根据ResNet版本调整输入维度）
    model.fc = nn.Sequential(
        nn.Linear(512, 256),
        nn.ReLU(),
        nn.Linear(256, 256),
        nn.ReLU(),
        nn.Linear(256, 3)
    )

    model = model.to(device)

    # 后续训练代码...
    # ...
    return model

说明 #

• 工具函数insert_dropout_after_conv会遍历模型所有层级的模块，确保每个Conv2d层后都插入Dropout。
• 若使用ResNet50等大模型，需将全连接层的输入维度从512改为2048，匹配对应模型的特征输出维度。

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