PyTorch模型层冻结:迁移学习中的参数固定原理与工程实践

PyTorch模型层冻结:迁移学习中的参数固定原理与工程实践 1. 项目概述为什么我们需要“冻结”网络层在深度学习的模型训练与微调实践中我们经常会遇到一个非常实际的需求只更新模型的一部分参数而让另一部分参数保持“冻结”状态不参与反向传播和梯度更新。这个操作通常被称为“冻结”Freezing或“固定”Fixing网络层。你可能会好奇模型训练不就是为了让所有参数都学习到最佳值吗为什么还要“偷懒”只更新一部分呢这背后有几个核心的驱动力。最常见的就是迁移学习。当我们拿到一个在ImageNet上预训练好的ResNet、BERT或GPT模型时它已经具备了强大的特征提取能力。如果我们想将其应用到一个新的、但数据量较小的特定任务比如医学影像分类或特定领域的文本情感分析上最明智的做法不是从头训练而是利用这个预训练好的模型作为起点。此时我们希望模型底层的通用特征提取器例如CNN的前几层或Transformer的前几层保持稳定因为它们已经学会了识别边缘、纹理、基础语法等通用模式。我们只需要针对新任务微调模型顶部的、更任务相关的层如分类头。冻结底层可以防止在小数据集上过拟合并大幅加快训练速度。另一个场景是多任务学习或分阶段训练。在复杂的模型架构中我们可能希望先训练好某个子模块比如一个特征编码器在后续训练其他部分时这个子模块的参数应该保持不变。又或者当我们向已有模型中插入新的层如适配器Adapter时通常也会选择冻结原始模型只训练新插入的部分以保证原始能力不丢失。在PyTorch中实现层的固定本质上就是操控张量的requires_grad属性并巧妙地管理优化器。听起来简单但实操中却有不少坑比如冻结后模型评估模式eval()的设置、批量归一化层BatchNorm在冻结时的特殊行为、以及如何优雅地实现部分层解冻进行分阶段训练等。接下来我将结合多年调参经验为你拆解在PyTorch中固定网络层的核心原理、多种实现方案以及那些官方文档里不会写的避坑指南。2. 核心原理requires_grad与优化器的双簧戏要理解如何固定层必须深入PyTorch的自动微分机制。关键就在于张量的requires_grad属性。2.1requires_grad梯度计算的开关在PyTorch中每个torch.Tensor都有一个布尔类型的requires_grad属性。当它被设置为True时PyTorch会在前向传播过程中追踪所有针对该张量的操作并构建一个动态计算图。在反向传播时会根据这个计算图计算该张量相对于某个标量损失函数的梯度。如果requires_gradFalse则该张量不会出现在计算图中其梯度既不会被计算自然也不会在反向传播中被更新。对于模型参数通常是nn.Parameter它是Tensor的子类这个属性决定了它是否参与学习。当我们从nn.Module中获取参数时例如通过model.parameters()或model.named_parameters()得到的每个参数都带有这个属性。注意直接修改模型层如model.conv1.weight.requires_grad False是有效的但更规范的做法是通过遍历模块的参数来设置因为一个模块可能包含多个参数张量。2.2 优化器只更新需要梯度的参数优化器如torch.optim.SGD,Adam在初始化时会接收一个需要优化的参数迭代器。在每一步optimizer.step()中优化器会遍历它管理的所有参数并根据其梯度存储在param.grad中更新参数值。这里有一个至关重要的细节优化器只关心在初始化时传给它的那些参数。如果你在创建优化器之后才修改某些参数的requires_grad属性优化器仍然会尝试更新这些参数如果它们的梯度不为None但这通常不是我们想要的行为。因此一个最佳实践是先设置好所有参数的requires_grad属性然后再用过滤后的参数列表来初始化优化器。2.3 计算图与内存效率将参数的requires_grad设置为False还有一个额外好处节省计算资源和内存。在前向传播中PyTorch无需为这些参数记录中间变量的计算历史这减少了构建计算图的开销。在反向传播时由于梯度计算被跳过也能节省显存和计算时间。这对于微调大型模型如ViT、Swin Transformer尤为重要。3. 实操方案三种主流冻结方法详解理论清晰后我们来看具体怎么做。我将介绍三种从基础到进阶的冻结方法并分析各自的适用场景。3.1 方法一遍历参数手动设置requires_grad这是最直接、最基础的方法。思路是遍历模型的所有参数根据参数名或所属层来判断是否需要冻结。import torch import torch.nn as nn import torchvision.models as models # 1. 加载预训练模型 model models.resnet50(pretrainedTrue) # 2. 冻结所有参数通常不是最终目的先全部冻结再解冻部分更常见 for param in model.parameters(): param.requires_grad False # 3. 解冻最后一层fc层的参数用于微调 for param in model.fc.parameters(): param.requires_grad True # 4. 仅将 requires_gradTrue 的参数传递给优化器 optimizer torch.optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), lr1e-3 )场景与技巧按层名冻结如果你要冻结所有卷积层但保留全连接层可以这样遍历for name, param in model.named_parameters(): if conv in name: # 根据参数名包含的关键字判断 param.requires_grad False先冻再解我个人的习惯是在微调时先for param in model.parameters(): param.requires_grad False全部冻结然后再针对性地解冻需要训练的层。这样逻辑更清晰不易出错。注意BN层对于包含批量归一化BatchNorm层的模型如ResNet需要特别注意。BN层在训练和评估时的行为不同。即使冻结了它的权重weight和偏置bias它在训练时仍然会累积运行均值running_mean和方差running_var。如果你希望BN层完全保持预训练状态不进行任何更新除了冻结参数还需要将其设置为评估模式model.bn1.eval()。但这样做在微调时可能会影响性能通常建议对于微调任务BN层的参数weight/bias可以保持可训练状态或者使用较小的学习率。3.2 方法二利用torch.no_grad上下文管理器torch.no_grad()是一个上下文管理器它会在其作用域内将所有张量操作的requires_grad计算强制设置为False。它通常用于模型推理评估阶段以节省内存。但我们也可以巧妙地用它来“临时冻结”部分网络的前向传播。class PartiallyFrozenModel(nn.Module): def __init__(self, backbone, head): super().__init__() self.backbone backbone self.head head # 冻结backbone的所有参数 for param in self.backbone.parameters(): param.requires_grad False def forward(self, x): # 在 no_grad 环境下运行 backbone确保其计算不被追踪 with torch.no_grad(): features self.backbone(x) # head 部分正常计算梯度可以回传 output self.head(features) return output # 使用示例 backbone models.resnet50(pretrainedTrue) head nn.Linear(backbone.fc.in_features, 10) # 新的分类头 model PartiallyFrozenModel(backbone, head) # 此时只有 head 的参数需要梯度 optimizer torch.optim.Adam(model.head.parameters(), lr1e-3)适用场景与局限 这种方法将冻结逻辑明确写在了forward函数里结构清晰。它特别适合** backbone 和 head 分离非常明确**的架构。但是它有一个潜在问题由于backbone在no_grad()下运行从backbone到head的路径在计算图上是断开的。这意味着如果你希望head的梯度能够以某种方式影响backbone的中间层在某些特殊设计下这是做不到的。对于标准的微调只更新head这完全没问题。3.3 方法三创建自定义优化器过滤参数这是最灵活、也是最推荐的方法尤其当冻结模式比较复杂时。核心思想是我们不对模型本身的requires_grad做过多改动而是通过构建不同的参数组Parameter Groups来告诉优化器应该更新哪些参数。import torch import torch.nn as nn import torchvision.models as models model models.resnet50(pretrainedTrue) # 定义需要训练的参数列表 params_to_update [] frozen_params [] for name, param in model.named_parameters(): # 例如我们冻结所有 stage1 到 stage3 的层只训练 stage4 和 fc if layer4 in name or fc in name: params_to_update.append(param) print(f训练层: {name}) else: frozen_params.append(param) param.requires_grad False # 可选用于节省计算图内存 print(f冻结层: {name}) # 关键优化器只接收需要训练的参数 optimizer torch.optim.SGD(params_to_update, lr0.001, momentum0.9) # 或者更优雅地使用参数组可以为不同层设置不同学习率 optimizer torch.optim.SGD([ {params: frozen_params, lr: 0}, # 学习率为0相当于不更新 {params: model.layer4.parameters(), lr: 1e-4}, # 浅层微调小学习率 {params: model.fc.parameters(), lr: 1e-3}, # 分类头大学习率 ], momentum0.9)优势灵活性高可以轻松实现不同层不同学习率差分学习率这是微调中的高级技巧。通常越靠近输出的层学习率可以设置得越大。逻辑清晰模型定义和训练逻辑分离。模型结构保持“完整”训练策略哪些层训练、以多大速率训练由优化器配置决定。便于调试通过打印named_parameters和检查requires_grad状态可以一目了然地确认冻结是否生效。4. 避坑指南与高级技巧在实际操作中仅仅设置requires_grad可能还不够下面这些经验之谈能帮你避免很多深夜调试的烦恼。4.1 批量归一化BatchNorm层的处理这是冻结操作中最容易踩坑的地方。BN层在训练model.train()和评估model.eval()模式下行为迥异。训练模式使用当前批次的统计量均值/方差进行归一化并更新其内部的运行统计量running_mean/ running_var。评估模式使用训练阶段累积得到的运行统计量进行归一化不再更新。问题当你冻结了BN层的weight和bias但模型仍处于train()模式时它的running_mean和running_var仍然会随着数据批次而更新这可能导致模型行为在训练过程中缓慢漂移。解决方案彻底冻结BN如果你确信预训练数据集的BN统计量非常适合你的新数据可以完全冻结BN层。for name, module in model.named_modules(): if isinstance(module, nn.BatchNorm2d): # 对于1D/3D BN同理 module.eval() # 设置为评估模式停止running stat更新 for param in module.parameters(): param.requires_grad False # 冻结可学习参数注意在训练循环中如果用了model.train()会将所有模块重置为训练模式因此需要在每个epoch或每个batch的训练步骤前重新将BN层设置为eval()。更稳妥的做法是写一个自定义的train()方法。微调BN更常见许多研究表明在领域迁移任务中微调BN层的参数和统计量能带来性能提升。此时可以保持BN层可训练但可能使用更小的学习率。# 在优化器参数组中为BN层单独设置更小的学习率 bn_params [p for n, p in model.named_parameters() if bn in n or norm in n] other_params [p for n, p in model.named_parameters() if not (bn in n or norm in n)] optimizer torch.optim.Adam([ {params: bn_params, lr: 1e-5}, {params: other_params, lr: 1e-3} ])4.2 验证冻结是否生效设置完后如何确认你的冻结操作真的起作用了我常用的调试检查点如下打印参数状态for name, param in model.named_parameters(): print(f{name}: requires_grad{param.requires_grad})查看目标层的requires_grad是否为False。检查梯度在完成一次loss.backward()之后检查冻结层的梯度是否为None。loss.backward() for name, param in model.named_parameters(): if param.requires_grad False: print(f{name} grad is: {param.grad}) # 应该输出 None观察参数变化在训练几个批次后直接比较冻结层参数的数值是否发生变化。initial_weight model.conv1.weight.clone().detach() # 克隆初始值 # ... 进行若干步训练 ... final_weight model.conv1.weight print(torch.equal(initial_weight, final_weight)) # 冻结了应该为 True4.3 分阶段解冻与渐进式微调对于某些任务更优的策略不是一开始就固定好哪些层训练哪些层冻结而是采用渐进式解冻。策略先只训练最顶部的几层如分类头。训练几个epoch后解冻相邻的下一部分层如ResNet的layer4同时可能降低整体学习率。继续训练再解冻更深的层如layer3以此类推。好处让模型逐步适应新数据避免一开始就剧烈更新深层特征导致预训练知识被“冲掉”。这在目标数据集与预训练数据集差异较大时特别有效。实现这需要你在训练循环的特定epoch动态地修改优化器中的参数组。你需要重新初始化优化器或者使用更高级的优化器封装如torch.optim.lr_scheduler.LambdaLR配合自定义函数来实现参数组学习率的变化。4.4 使用torch.compile与新特性如果你使用较新版本的PyTorch2.0并使用了torch.compile来加速模型需要注意编译后的图优化可能会对梯度计算产生影响。确保在编译前就完成参数的冻结设置。一般来说requires_grad的设置与torch.compile是兼容的但作为最佳实践建议的代码顺序是定义模型 - 设置冻结 - 编译模型 - 定义优化器。5. 完整代码示例一个可复现的微调流程最后我将给出一个完整的、包含数据加载、模型冻结、训练和验证的示例以ResNet50在CIFAR-10上的微调为例。import torch import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as transforms from torchvision import models # 1. 设备设置 device torch.device(cuda if torch.cuda.is_available() else cpu) # 2. 数据加载 (CIFAR-10) transform_train transforms.Compose([ transforms.RandomCrop(32, padding4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) transform_test transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) trainset torchvision.datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform_train) trainloader torch.utils.data.DataLoader(trainset, batch_size128, shuffleTrue, num_workers2) testset torchvision.datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformtransform_test) testloader torch.utils.data.DataLoader(testset, batch_size100, shuffleFalse, num_workers2) # 3. 模型加载与修改 model models.resnet50(pretrainedTrue) # 修改最后的全连接层CIFAR-10是10类 num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, 10) # 4. 冻结除最后一层外的所有参数 for name, param in model.named_parameters(): if fc not in name: # 只保留全连接层可训练 param.requires_grad False else: param.requires_grad True model model.to(device) # 5. 优化器与损失函数 (只优化 requires_gradTrue 的参数) optimizer optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr0.001, momentum0.9, weight_decay5e-4) criterion nn.CrossEntropyLoss() scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max200) # 6. 训练与验证循环 def train(epoch): model.train() running_loss 0.0 for batch_idx, (inputs, targets) in enumerate(trainloader): inputs, targets inputs.to(device), targets.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, targets) loss.backward() optimizer.step() running_loss loss.item() scheduler.step() print(fEpoch: {epoch}, Loss: {running_loss/len(trainloader):.4f}) def test(): model.eval() correct 0 total 0 with torch.no_grad(): for inputs, targets in testloader: inputs, targets inputs.to(device), targets.to(device) outputs model(inputs) _, predicted outputs.max(1) total targets.size(0) correct predicted.eq(targets).sum().item() accuracy 100. * correct / total print(fTest Accuracy: {accuracy:.2f}%) return accuracy # 7. 开始训练 for epoch in range(1, 11): # 训练10个epoch train(epoch) test()这个流程清晰地展示了从模型准备、参数冻结、优化器配置到训练验证的完整步骤。你可以在此基础上尝试应用前面提到的BN层处理、差分学习率或渐进式解冻等高级技巧以适应更复杂的任务需求。记住没有一成不变的冻结策略最好的方法总是需要通过实验根据你的具体数据、模型和任务来验证和调整。