如何在PyTorch Lightning中训练RegNet-800MF骨干网络最后几层

你当前代码存在两个会导致微调失效的问题：一是全局调用freeze()冻结了骨干所有参数，二是forward中用torch.no_grad()包裹骨干前向过程、强制骨干处于eval()模式，就算修改参数的requires_grad属性也无法回传梯度、更新层权重。
具体实现按以下步骤修改即可：

1. 初始化阶段：先全冻骨干，再针对性开放最后几层梯度 #

RegNet-Y-800MF的特征提取结构从浅到深为stem→stage1→stage2→stage3→stage4→avgpool，新手建议先只解冻最深的stage4和avgpool，参数量占比低，不容易过拟合，效果达不到预期再考虑追加解冻stage3。
修改MechClassifier的初始化逻辑：

# 加载预训练检查点后，先全量冻结骨干参数
self.backbone.freeze()
self.backbone.eval()

# 解冻指定最后几层
unfreeze_layers = [
    self.backbone.model.stage4,
    self.backbone.model.avgpool
]
for layer in unfreeze_layers:
    for param in layer.parameters():
        param.requires_grad = True
    layer.train()

2. 删除forward中阻碍梯度回传的代码 #

去掉原forward里强制骨干为eval模式、以及torch.no_grad()的上下文包裹，未解冻的层因为requires_grad=False，PyTorch会自动跳过梯度计算，不会额外占用过多显存。修改后的forward如下：

def forward(self, x):
    x = self.backbone.model(x)
    out1 = self.mf_classifier(x)
    out2 = self.wd_classifier(x)
    return (out1, out2)

3. 添加训练周期钩子，固定冻结层的状态 #

为了防止训练时模型整体切train()模式，导致冻结层的BatchNorm统计量被污染，添加训练轮次启动钩子，单独控制每层状态：

def on_train_epoch_start(self):
    # 冻结层始终保持eval模式，不更新BN统计量
    self.backbone.model.stem.eval()
    self.backbone.model.stage1.eval()
    self.backbone.model.stage2.eval()
    self.backbone.model.stage3.eval()
    # 解冻层保持train模式，正常更新参数和BN统计量
    self.backbone.model.stage4.train()
    self.backbone.model.avgpool.train()

如果想让训练更稳定，可以加热身逻辑：前几轮完全冻结骨干只训新添加的分类头，几轮后再解冻最后几层联合训练，参考代码：

def on_train_epoch_start(self):
    # 前3轮全冻骨干，只训分类头
    if self.current_epoch < 3:
        self.backbone.freeze()
        self.backbone.eval()
    else:
        # 3轮后解冻最后几层
        self.backbone.model.stem.eval()
        self.backbone.model.stage1.eval()
        self.backbone.model.stage2.eval()
        self.backbone.model.stage3.eval()
        for param in self.backbone.model.stage4.parameters():
            param.requires_grad = True
        for param in self.backbone.model.avgpool.parameters():
            param.requires_grad = True
        self.backbone.model.stage4.train()
        self.backbone.model.avgpool.train()

4. 优化器分组配置参数（必做） #

微调时解冻的骨干层学习率需要比新加的分类头低（通常设为分类头学习率的1/10），避免预训练好的特征提取能力被大学习率破坏。重写configure_optimizers方法做参数分组：

def configure_optimizers(self):
    optimizer = torch.optim.AdamW([
        # 新分类头用初始学习率
        {"params": self.mf_classifier.parameters(), "lr": self.lr},
        {"params": self.wd_classifier.parameters(), "lr": self.lr},
        # 解冻的骨干层用1/10学习率
        {"params": list(self.backbone.model.stage4.parameters()) + list(self.backbone.model.avgpool.parameters()), "lr": self.lr * 0.1},
    ], weight_decay=self.weight_decay)
    return optimizer

新手注意事项 #

• 不要一开始就解冻过半的骨干层，小数据集下极容易过拟合，且训练显存、速度都会明显劣化
• 训练时注意观察训练集和验证集的精度差，如果差距过大说明过拟合，要减少解冻层数、加正则或者调小骨干学习率
• 不需要在代码里手动过滤requires_grad=False的参数，PyTorch优化器会自动跳过这部分参数不更新

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