PyTorch冻结神经网络层:原理、实战与进阶技巧
1. 项目概述为什么需要“冻结”神经网络层在深度学习的实际项目里尤其是做迁移学习或者微调预训练模型时我们经常会遇到一个非常具体的需求让模型的一部分参数“静默”下来不参与反向传播的梯度计算和更新。这个操作我们通常称之为“冻结”Freeze。听起来有点抽象我打个比方你拿到一个别人训练好的、非常厉害的厨师预训练模型他擅长做川菜。现在你想让他学做粤菜但希望他保留做川菜的看家本领比如对麻辣调味的精准把握只调整他处理海鲜和煲汤的技巧。这时候你就需要“冻结”他大脑中关于川菜调味的神经回路只让与海鲜处理相关的部分去学习和适应新任务。PyTorch作为当前主流的深度学习框架之一其动态计算图机制给了我们极大的灵活性但如何精准、高效地控制哪些层参与训练哪些层不参与却是一个需要明确技巧的实操点。直接设置requires_gradFalse是最核心的原理但实际应用中从优化器配置、BatchNorm层处理到多阶段训练策略里面有不少坑。新手常犯的错误是以为设置了requires_gradFalse就万事大吉结果发现模型不收敛或者性能诡异问题往往就出在细节没处理好。这篇文章我就结合自己多次在CV和NLP项目中微调ResNet、BERT等模型的经验把PyTorch中冻结层的各种方法、背后的原理、常见的“坑”以及高级技巧给你一次讲透。2. 核心原理与基础操作理解requires_grad与优化器在深入具体代码之前我们必须先夯实两个核心概念张量的requires_grad属性和优化器Optimizer的工作机制。这是理解“冻结”为何生效的基石。2.1requires_grad梯度计算的开关在PyTorch中每个torch.Tensor都有一个布尔类型的属性requires_grad。当它被设置为True时PyTorch会在前向传播过程中追踪所有针对该张量的操作并构建一个动态计算图。在反向传播时会根据这个计算图计算该张量通常是模型参数的梯度。如果requires_gradFalse则该张量不会出现在计算图中其梯度既不会被计算自然也不会被更新。模型中的参数model.parameters()本质上就是requires_gradTrue的张量。因此冻结某一层最直接的方法就是将该层所有参数的requires_grad属性设置为False。import torch import torch.nn as nn # 以一个简单的两层网络为例 class SimpleNet(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(10, 20) self.fc2 nn.Linear(20, 1) def forward(self, x): x self.fc1(x) x self.fc2(x) return x model SimpleNet() # 打印初始状态 for name, param in model.named_parameters(): print(f{name}: requires_grad {param.requires_grad}) # 输出 # fc1.weight: requires_grad True # fc1.bias: requires_grad True # fc2.weight: requires_grad True # fc2.bias: requires_grad True # 冻结第一层 fc1 for param in model.fc1.parameters(): param.requires_grad False # 再次打印 for name, param in model.named_parameters(): print(f{name}: requires_grad {param.requires_grad}) # 输出 # fc1.weight: requires_grad False # fc1.bias: requires_grad False # fc2.weight: requires_grad True # fc2.bias: requires_grad True2.2 优化器只更新需要梯度的参数仅仅设置requires_gradFalse还不够。优化器如torch.optim.SGD或Adam在初始化时会传入一个需要优化的参数迭代器通常是model.parameters()。优化器内部会存储这些参数的引用。在每次optimizer.step()时优化器会遍历它存储的所有参数并根据其梯度.grad属性更新参数值。关键点来了如果一个参数的requires_gradFalse在反向传播中它的.grad属性会是None。然而优化器仍然会尝试“更新”它只不过因为梯度是None所以更新量是零相当于没变。但这会带来两个问题1不必要的计算开销2在某些复杂的优化器如Adam中它可能还会更新其内部状态如动量这可能导致意料之外的行为或错误。因此最佳实践是在创建优化器时只传入那些真正需要训练的参数。这通过filter函数可以轻松实现。import torch.optim as optim # 错误做法优化器仍然包含了所有参数包括被冻结的 # optimizer optim.Adam(model.parameters(), lr0.001) # 正确做法只优化 requires_gradTrue 的参数 trainable_params filter(lambda p: p.requires_grad, model.parameters()) optimizer optim.Adam(trainable_params, lr0.001)注意这是一个极易被忽略但至关重要的步骤。很多同学冻结了层但效果不佳排查半天才发现优化器没做过滤导致训练效率低下甚至出现奇怪问题。3. 实战演练冻结预训练模型的一部分现在我们用一个更贴近实际的场景来演练微调一个在ImageNet上预训练的ResNet-18模型用于一个具体的花卉分类任务。假设我们想冻结前面的卷积层特征提取器只训练最后的全连接分类头。3.1 模型准备与结构分析首先加载预训练模型并查看其结构确定我们要冻结哪些部分。import torchvision.models as models # 加载预训练的ResNet-18 model models.resnet18(pretrainedTrue) # 查看模型的主要组成部分 print(model) # 输出会显示类似结构 # ResNet( # (conv1): Conv2d(3, 64, kernel_size(7, 7), stride(2, 2), padding(3, 3), biasFalse) # (bn1): BatchNorm2d(64, eps1e-05, momentum0.1, affineTrue, track_running_statsTrue) # (relu): ReLU(inplaceTrue) # (maxpool): MaxPool2d(kernel_size3, stride2, padding1, dilation1, ceil_modeFalse) # (layer1): Sequential(...) # (layer2): Sequential(...) # (layer3): Sequential(...) # (layer4): Sequential(...) # (avgpool): AdaptiveAvgPool2d(output_size(1, 1)) # (fc): Linear(in_features512, out_features1000, biasTrue) # )对于ResNetconv1、bn1、layer1到layer4都是卷积层用于提取特征。fc是最后的全连接层输出1000维对应ImageNet的1000个类别。我们的任务是将其改为输出花卉类别数比如5类。3.2 冻结所有卷积层我们的策略是冻结除了最后一层fc之外的所有参数。# 冻结除最后一层全连接层外的所有参数 for name, param in model.named_parameters(): # 如果参数名不以‘fc.’开头则冻结它 if not name.startswith(fc.): param.requires_grad False # 验证冻结效果 for name, param in model.named_parameters(): if param.requires_grad: print(f可训练参数: {name}) else: print(f已冻结参数: {name}) # 修改最后一层全连接层以适应新的分类任务假设有5类花卉 num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, 5) # 新的fc层默认 requires_gradTrue # 再次确认新的fc层是可训练的其他层是冻结的 print(\n修改fc层后) for name, param in model.named_parameters(): if fc in name: print(f{name}: requires_grad {param.requires_grad})3.3 配置优化器与训练循环根据可训练参数配置优化器并编写训练循环。import torch.optim as optim from torch.utils.data import DataLoader # 假设 train_loader 是你的训练数据加载器 # 只对需要梯度的参数进行优化 optimizer optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr0.001) criterion nn.CrossEntropyLoss() model.train() for epoch in range(num_epochs): for inputs, labels in train_loader: optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() # ... 其他日志记录等代码在这个循环中反向传播loss.backward()只会为requires_gradTrue的参数计算梯度。由于优化器只包含了这些参数optimizer.step()只会更新它们。被冻结的层其参数值在整个训练过程中将保持不变。4. 进阶技巧与疑难杂症处理基础操作看似简单但在复杂的模型和训练策略下会遇到一些棘手的问题。下面分享几个我踩过坑后总结的进阶技巧。4.1 处理BatchNorm层冻结还是解冻这是一个非常经典的问题。BatchNormBN层有可学习的参数weight和bias和运行时统计量running_mean和running_var。当我们冻结其所在的卷积块时BN层该如何处理情景一完全冻结特征提取器。如果你希望特征提取器完全保持预训练时的状态那么BN层的参数和运行时统计量都应该冻结。这意味着设置param.requires_grad False并且在训练时将其设置为eval()模式以防止运行时统计量更新。# 冻结所有参数并将模型设为eval模式以防止BN统计量更新 for param in model.parameters(): param.requires_grad False model.eval() # 但注意这样你就无法训练任何层了。通常我们只冻结部分层。 # 更常见的做法是遍历模型将需要冻结的模块如layer1设为eval() model.layer1.eval() for param in model.layer1.parameters(): param.requires_grad False然而更简洁的做法是利用torch.no_grad()上下文管理器但这更适合推理阶段。对于训练中的部分冻结手动设置eval()和requires_grad更可靠。情景二微调特征提取器但学习率很低。此时BN层的参数应该随着卷积层一起以极低的学习率进行微调。运行时统计量则继续在训练中更新。这是最常见和推荐的做法因为BN层的统计量对输入分布敏感新数据上更新统计量通常是有益的。实操心得对于大多数计算机视觉的微调任务我建议不要单独冻结BN层。将BN层与其前面的卷积层视为一个整体采用相同的处理策略冻结或使用极小的学习率微调。单独冻结BN层参数但更新其统计量容易导致内部协变量偏移使训练不稳定。4.2 差异化学习率Differential Learning Rates这是微调中的高级技巧也是冻结策略的灵活变体。我们不是简单地将某些层的学习率设为0冻结而是给不同层分配不同的学习率。通常靠近输入的层使用更小的学习率因为它们提取通用特征靠近输出的层特别是新添加的层使用更大的学习率。PyTorch可以通过优化器参数组param_groups来实现。# 假设我们想将模型参数分为三组 # 1. 卷积层layer1, layer2的参数学习率很小如1e-5 # 2. 卷积层layer3, layer4的参数学习率中等如1e-4 # 3. 全连接层fc的参数学习率较大如1e-3 ignored_params [] conv1_params [] conv2_params [] fc_params [] for name, param in model.named_parameters(): if layer1 in name or layer2 in name: conv1_params.append(param) elif layer3 in name or layer4 in name: conv2_params.append(param) elif fc in name: fc_params.append(param) else: # 其他层如conv1, bn1等可以根据需要归类 ignored_params.append(param) optimizer optim.Adam([ {params: conv1_params, lr: 1e-5}, {params: conv2_params, lr: 1e-4}, {params: fc_params, lr: 1e-3}, ], lr0.001) # 这里的lr作为默认值会被各组的lr覆盖这样我们就实现了比简单冻结更精细的控制。4.3 多阶段训练与动态解冻有时我们可能采用“热身”或“渐进解冻”策略。例如第一阶段只训练新添加的分类头冻结所有骨干网络。第二阶段解冻骨干网络的最后1-2个阶段如layer4以较小的学习率进行微调。第三阶段解冻更多层继续微调。这需要我们在训练过程中动态地修改参数的requires_grad属性并重新创建优化器。因为优化器在初始化时就固定了要优化的参数集合。# 第一阶段只训练fc层 for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True optimizer optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr0.01) # ... 训练若干轮 ... # 第二阶段解冻layer4并重新创建优化器 for param in model.layer4.parameters(): param.requires_grad True # 必须重新创建优化器以包含新解冻的参数 optimizer optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr0.001) # ... 继续训练 ...重要提示动态改变requires_grad后务必重新创建优化器。旧的优化器内部状态如动量不会自动适配新加入的参数继续使用会导致错误。5. 常见问题排查与调试技巧即使按照上述步骤操作你可能还是会遇到一些问题。这里列几个我常被问到的情况和排查思路。5.1 问题模型似乎完全没在学习损失不下降排查点1优化器参数过滤。这是最常见的原因。检查你是否在创建优化器时使用了filter(lambda p: p.requires_grad, model.parameters())。打印优化器参数组的长度来验证。optimizer optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr0.001) print(f“优化器参数数量 {len(optimizer.param_groups[0][‘params’])}”) # 对比一下总参数数量 total_params sum(p.numel() for p in model.parameters()) trainable_params sum(p.numel() for p in model.parameters() if p.requires_grad) print(f“总参数 {total_params}, 可训练参数 {trainable_params}”)排查点2梯度检查。在训练循环中检查关键层的梯度是否为None。loss.backward() # 检查某个冻结层的梯度 print(model.conv1.weight.grad) # 应该为 None # 检查某个可训练层的梯度 print(model.fc.weight.grad) # 应该是一个Tensor排查点3学习率是否过低。对于微调特别是只训练最后几层时学习率不宜过小。可以尝试从1e-3或1e-4开始。5.2 问题训练过程中出现NaN损失或梯度爆炸排查点1BatchNorm层模式。如果你冻结了包含BN层的模块但没有将其设置为eval()模式在训练时它的running_mean/var仍在更新而前方的卷积层已被冻结可能导致数据分布变化异常。尝试将冻结的模块设为eval()。model.layer1.eval() # 冻结的层 model.layer2.train() # 可训练的层 # 注意model.train() 会将所有子模块设为train模式会覆盖上面的eval设置。 # 因此需要在每个epoch或batch的训练循环开始前显式设置各模块的模式。排查点2差异化学习率设置过激。如果不同层的学习率差异巨大如相差1000倍可能导致优化过程不稳定。适当缩小学习率差距或使用学习率预热Warmup策略。5.3 问题如何验证层确实被冻结了最直观的方法是监控参数值的变化。可以在训练前后保存特定层的参数并进行比较。# 训练前保存fc层和某个卷积层的权重 fc_weight_before model.fc.weight.data.clone() conv_weight_before model.conv1.weight.data.clone() # ... 进行一个epoch的训练 ... fc_weight_after model.fc.weight.data conv_weight_after model.conv1.weight.data # 检查变化 print(“FC层权重是否变化”, not torch.allclose(fc_weight_before, fc_weight_after)) print(“Conv1层权重是否变化”, not torch.allclose(conv_weight_before, conv_weight_after)) # 期望输出FC层 True Conv1层 False6. 封装与最佳实践编写可复用的冻结函数为了提升代码的整洁度和复用性可以编写一些辅助函数。def set_parameter_requires_grad(model, layer_names, requires_gradFalse): “”” 设置模型中指定层名的参数的 requires_grad 属性。 Args: model: nn.Module layer_names: list of str, 层名支持部分匹配如 ‘layer1’ requires_grad: bool “”” for name, param in model.named_parameters(): for layer_name in layer_names: if layer_name in name: param.requires_grad requires_grad break # 匹配到一个就跳出内层循环 def get_trainable_parameters(model): “””获取模型中所有可训练参数””” return filter(lambda p: p.requires_grad, model.parameters()) def freeze_all_except(model, except_layer_names): “”” 冻结除指定层之外的所有参数。 Args: model: nn.Module except_layer_names: list of str, 不解冻的层名列表 “”” for param in model.parameters(): param.requires_grad False for name, param in model.named_parameters(): for layer_name in except_layer_names: if layer_name in name: param.requires_grad True break # 使用示例 model models.resnet18(pretrainedTrue) # 方法1冻结除了 ‘fc’ 和 ‘layer4’ 以外的所有层 freeze_all_except(model, [‘fc’, ‘layer4’]) # 方法2更精细地控制只冻结 ‘layer1’ 和 ‘layer2’ set_parameter_requires_grad(model, [‘layer1’, ‘layer2’], requires_gradFalse) optimizer optim.Adam(get_trainable_parameters(model), lr0.001)最后关于是否要冻结、冻结多少层并没有绝对的金科玉律。这取决于你的数据集大小、与预训练数据集的相似度以及任务本身。一个实用的启发性原则是数据量越小、任务与预训练任务越不相关就该冻结越多的层以避免过拟合。反之则可以解冻更多的层进行微调。在实际项目中我通常会设计一个简单的实验尝试不同的冻结策略例如冻结全部卷积层、冻结前三个stage、只冻结第一个stage等在验证集上比较性能选择最优的方案。模型训练本身就是一个需要不断实验和调整的过程理解工具背后的原理才能让你在调参时更有方向少走弯路。