PyTorch Lightning 迁移学习与微调实战从冻结特征提取器到 BackboneFinetuning 自动微调【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000 GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning本文基于 PyTorch Lightning 官方文档《Transfer Learning》见 docs/source-pytorch/advanced/finetuning.rst其通过.. include::引入 transfer_learning.rst结合仓库源码深入讲解如何在 Lightning 中加载预训练模型完成迁移学习Transfer Learning与微调Finetuning。本指南面向希望借助预训练模型加速训练的开发者涵盖三类核心场景把任意 PyTorchnn.Module接入 Lightning、把预训练的LightningModule当作特征提取器复用、以及使用官方BackboneFinetuning/BaseFinetuning回调实现先冻结主干、再逐步解冻的自动化微调流程。读完本文你将能够把 ImageNet 预训练 ResNet、HuggingFace BERT 等模型无缝嵌入自己的 Lightning 项目并掌握学习率分层调度、参数组管理与断点续训等实战细节。一、前提任何 PyTorch nn.Module 都能与 Lightning 协作Lightning 对迁移学习的核心承诺是只要是torch.nn.Module子类就可以直接用在 Lightning 训练流程中。原因很直接——LightningModule本身就是nn.Module的子类因此你可以把任意预训练nn.Module如 torchvision 的 ResNet、HuggingFace 的 BERT作为属性嵌入LightningModule也可以把已经训练好的LightningModule整体或其中一部分如 encoder拿出来作为另一个LightningModule的子模块。这一设计保证了 Lightning 对迁移学习框架的选择完全无感torchvision、transformers、timm 等生态的模型都能直接复用无需任何包装代码。二、复用预训练的 LightningModule把 AutoEncoder 变成特征提取器文档中的第一个完整示例演示了如何把预训练的AutoEncoder本身是一个LightningModule作为特征提取器嵌入到一个新的分类模型中class Encoder(torch.nn.Module): ... class AutoEncoder(LightningModule): def __init__(self): self.encoder Encoder() self.decoder Decoder() class CIFAR10Classifier(LightningModule): def __init__(self): # init the pretrained LightningModule self.feature_extractor AutoEncoder.load_from_checkpoint(PATH).encoder self.feature_extractor.freeze() # the autoencoder outputs a 100-dim representation and CIFAR-10 has 10 classes self.classifier nn.Linear(100, 10) def forward(self, x): representations self.feature_extractor(x) x self.classifier(representations) ...这段代码体现了两条关键 APIload_from_checkpoint(PATH)这是 Lightning 官方推荐的加载方式。它是类方法在 src/lightning/pytorch/core/module.py 中定义Lightning 保存 checkpoint 时会同时存储__init__传入的超参数存于hyper_parameters键下因此加载时可以直接恢复模型结构与权重。它还支持map_locationGPU/CPU 间映射、hparams_file单独提供超参 yaml、strict是否严格匹配 state_dict 键、weights_only从不可信来源加载时的安全选项以及**kwargs覆盖保存时的超参数。需要注意必须用 LightningModule 的类来调用用实例调用会抛出TypeError。freeze()LightningModule 的内置方法module.py将所有参数的requires_grad置为False并将模型切到eval()模式返回自身便于链式调用。对应地unfreeze()module.py恢复所有参数可训练并切回train()模式。在分类器场景中self.feature_extractor.freeze()确保反向传播时只更新classifier的参数大幅降低显存与计算开销。这样我们就用预训练的 AutoEncoder 完成了编码器 分类头的迁移学习组合。需要特别提醒由于load_from_checkpoint是类方法上述写法AutoEncoder.load_from_checkpoint(PATH)返回的正是AutoEncoder实例取.encoder属性即可获得纯 PyTorch 编码器模块。三、实战示例ImageNet 预训练模型在 CIFAR-10 上的微调计算机视觉文档给出了第二个完整实战用 torchvision 提供的 ImageNet 预训练 ResNet50 提取特征在 CIFAR-10 上微调分类头。3.1 定义模型截断主干 新分类头import torchvision.models as models class ImagenetTransferLearning(LightningModule): def __init__(self): super().__init__() # init a pretrained resnet backbone models.resnet50(weightsDEFAULT) num_filters backbone.fc.in_features layers list(backbone.children())[:-1] self.feature_extractor nn.Sequential(*layers) self.feature_extractor.eval() # use the pretrained model to classify cifar-10 (10 image classes) num_target_classes 10 self.classifier nn.Linear(num_filters, num_target_classes) def forward(self, x): with torch.no_grad(): representations self.feature_extractor(x).flatten(1) x self.classifier(representations) ...要点拆解models.resnet50(weightsDEFAULT)加载 ImageNet 预训练权重backbone.fc.in_features读取原全连接层输入维度ResNet50 为 2048list(backbone.children())[:-1]去掉最后的fc分类层剩余部分作为通用特征提取器with torch.no_grad()包裹前向传播确保特征提取阶段不产生梯度结合手动freeze使用新的classifier是一个nn.Linear(num_filters, num_target_classes)输出目标类别数。3.2 微调与预测# 微调 model ImagenetTransferLearning() trainer Trainer() trainer.fit(model) # 预测 model ImagenetTransferLearning.load_from_checkpoint(PATH) model.freeze() x some_images_from_cifar10() predictions model(x)流程非常直接Trainer().fit(model)完成微调checkpoint 自动保存后用load_from_checkpoint(PATH)恢复权重再freeze()冻结全部参数后即可对自有数据推理。仓库中提供了更完整的 ImageNet 训练脚本 examples/pytorch/domain_templates/imagenet.py它通过LightningCLI暴露--model.data_path等参数并演示了ModelCheckpoint(monitorval_acc1, modemax)按验证集 top-1 精度保存最优模型的工程化做法可作为大规模训练时的参考模板。文档同时点明一个关键认知学术场景里我们在 ImageNet 预训练、CIFAR-10 上微调、再在 CIFAR-10 上预测而实际业务中等价的做法是在某个大数据集上预训练、在你的小数据集上微调、再在你的数据集上预测。这正是迁移学习的价值所在——用少量标注数据即可获得强于从头训练的效果。四、实战示例BERT 微调自然语言处理Lightning 对 NLP 迁移学习同样完全开放只要模型是torch.nn.Module子类即可。文档给出了基于 HuggingFace transformers 的 MNLI 微调示例class BertMNLIFinetuner(LightningModule): def __init__(self): super().__init__() self.bert BertModel.from_pretrained(bert-base-cased, output_attentionsTrue) self.bert.train() self.W nn.Linear(bert.config.hidden_size, 3) self.num_classes 3 def forward(self, input_ids, attention_mask, token_type_ids): h, _, attn self.bert(input_idsinput_ids, attention_maskattention_mask, token_type_idstoken_type_ids) h_cls h[:, 0] logits self.W(h_cls) return logits, attn模式与 CV 场景完全同构预训练 BERT 负责把 token 序列编码为上下文表示取[CLS]位置的输出h[:, 0]作为整句表示再通过一个线性层self.W映射到 3 个分类 logits。output_attentionsTrue使前向额外返回注意力矩阵可用于可解释性分析。在training_step/validation_step中按常规方式组织 batchinput_ids、attention_mask、token_type_ids即可接入 Trainer 训练。这再次印证迁移学习的接缝处只有一个——新任务的头其余全部复用预训练权重。五、自动化微调BackboneFinetuning 回调手动冻结/解冻需要自己写循环逻辑。PyTorch Lightning 提供了开箱即用的 BackboneFinetuning 回调定义于 src/lightning/pytorch/callbacks/finetuning.py其核心策略是训练开始时冻结主干backbone只训练任务头到达指定 epoch 后逐步解冻主干并按调度函数提升其学习率。这对于大型预训练模型尤其有价值——既能避免前期灾难性遗忘又能控制训练成本。5.1 模型结构要求回调要求模型具备特定结构class MyModel(LightningModule): def __init__(self): super().__init__() # REQUIRED: Your model must have a backbone attribute # This should be the pretrained part you want to finetune self.backbone some_pretrained_model # Your task-specific layers (head, classifier, etc.) self.head nn.Linear(backbone_features, num_classes) def configure_optimizers(self): # Only optimize the head initially - backbone will be added automatically return torch.optim.Adam(self.head.parameters(), lr1e-3)两条硬性约定在源码中均有对应校验必须有一个名为backbone的nn.Module属性。BackboneFinetuning.on_fit_startfinetuning.py会检查hasattr(pl_module, backbone) and isinstance(pl_module.backbone, Module)不满足则抛出MisconfigurationException(The LightningModule should have a nn.Module backbone attribute)。configure_optimizers只优化 head。backbone 由回调在解冻时自动通过optimizer.add_param_group加入无需也不应在初始化时传入。5.2 完整示例ResNet BackboneFinetuningimport torch import torch.nn as nn import torchvision.models as models from lightning.pytorch import LightningModule, Trainer from lightning.pytorch.callbacks import BackboneFinetuning class ResNetClassifier(LightningModule): def __init__(self, num_classes10, learning_rate1e-3): super().__init__() self.save_hyperparameters() # Create backbone from pretrained ResNet resnet models.resnet50(weightsDEFAULT) # Remove the final classification layer self.backbone nn.Sequential(*list(resnet.children())[:-1]) # Add custom classification head self.head nn.Sequential( nn.Flatten(), nn.Linear(resnet.fc.in_features, 512), nn.ReLU(), nn.Dropout(0.2), nn.Linear(512, num_classes) ) def forward(self, x): # Extract features with backbone features self.backbone(x) # Classify with head return self.head(features) def training_step(self, batch, batch_idx): x, y batch y_hat self(x) loss nn.functional.cross_entropy(y_hat, y) self.log(train_loss, loss) return loss def configure_optimizers(self): # Initially only train the head - backbone will be added by callback return torch.optim.Adam(self.head.parameters(), lrself.hparams.learning_rate) # Setup the finetuning callback backbone_finetuning BackboneFinetuning( unfreeze_backbone_at_epoch10, # Start unfreezing backbone at epoch 10 lambda_funclambda epoch: 1.5, # Gradually increase backbone learning rate backbone_initial_ratio_lr0.1, # Backbone starts at 10% of head learning rate should_alignTrue, # Align rates when backbone rate reaches head rate verboseTrue # Print learning rates during training ) model ResNetClassifier() trainer Trainer(callbacks[backbone_finetuning], max_epochs20)5.3 参数详解与源码一一对应BackboneFinetuning.__init__的完整签名finetuning.py如下各参数含义与默认值参数默认值含义unfreeze_backbone_at_epoch10解冻 backbone 的起始 epoch。在finetune_function中epoch unfreeze_backbone_at_epoch时执行首次解冻并创建 backbone 参数组epoch 该值时按lambda_func逐轮提升 backbone 学习率lambda_funcmultiplicative返回2.0学习率调度函数接收当前 epoch返回乘数每轮 backbone 新学习率 lambda_func(epoch 1) * previous_backbone_lrbackbone_initial_ratio_lr10e-2即 0.1用于按比例缩放 backbone 初始学习率初始 backbone lr current_lr * backbone_initial_ratio_lr当未显式指定backbone_initial_lr时backbone_initial_lrNone可选的 backbone 初始学习率绝对值优先于backbone_initial_ratio_lrshould_alignTrue当调度后的 backbone 学习率超过head 学习率时是否将 backbone 学习率钳制对齐到 head 学习率并保持至训练结束initial_denom_lr10.0首次解冻时若未提供 lr则用current_learning_rate / initial_denom_lr作为初始学习率unfreeze_and_add_param_group的参数train_bnTrue冻结时是否让 BatchNorm 层保持可训练影响freeze对 BN 的处理详见下文verboseFalse为True时每轮打印 head 与 backbone 的当前学习率精度由rounding控制rounding12打印学习率时保留的小数位数5.4 底层机制回调如何动优化器BackboneFinetuning继承自BaseFinetuningCallback的子类整个微调流程由两个钩子驱动setupfinetuning.py在configure_optimizers之前调用freeze_before_training(pl_module)即先把pl_module.backbone整体冻结。同时这里有一个重要限制若使用 DeepSpeed 策略会直接抛出NotImplementedError因为该回调的动态add_param_group与 DeepSpeed 的参数管理不兼容测试用例见 tests/tests_pytorch/callbacks/test_finetuning_callback.py。on_train_epoch_startfinetuning.py每个训练 epoch 开始时遍历trainer.optimizers调用finetune_function(pl_module, trainer.current_epoch, optimizer)执行解冻逻辑并调用_store把新增的 param_group 元数据以参数名而非张量引用保存便于 checkpoint 恢复记录到_internal_optimizer_metadata。首次解冻发生在epoch unfreeze_backbone_at_epoch时finetuning.py取optimizer.param_groups[0][lr]作为当前 head 学习率计算initial_backbone_lr然后调用unfreeze_and_add_param_group(pl_module.backbone, optimizer, initial_backbone_lr, ...)将 backbone 参数作为新的 param_group加入优化器。此后每轮next_current_backbone_lr self.lambda_func(epoch 1) * self.previous_backbone_lr next_current_backbone_lr ( current_lr if (self.should_align and next_current_backbone_lr current_lr) else next_current_backbone_lr ) optimizer.param_groups[-1][lr] next_current_backbone_lr即 backbone 学习率按lambda_func指数式增长一旦超过 head 学习率且should_alignTrue就对齐并保持。TestBackboneFinetuningCallback测试test_finetuning_callback.py专门断言了这种backbone lr 不超过 head lr、后期两者相等的收敛行为。关于 BatchNorm 需要特别说明BaseFinetuning.freezefinetuning.py在train_bnTrue时会通过make_trainable让 BN 层的weight/bias保持requires_gradTrue且track_running_statsTrue而train_bnFalse时则彻底冻结并关闭 running stats 更新。这一行为被test_freeze_unfreeze_functiontest_finetuning_callback.py完整验证。六、自定义微调策略继承 BaseFinetuning当BackboneFinetuning的内置调度不满足需求时可以继承 BaseFinetuning 自定义冻结/解冻策略。只需重写两个钩子文档原文如此freeze_before_training(pl_module)在configure_optimizers之前调用负责冻结任意模块的参数finetune_function(pl_module, epoch, optimizer)每个训练 epoch 开始时调用负责解冻参数并将其加入优化器的新 param_group。from lightning.pytorch.callbacks.finetuning import BaseFinetuning class CustomFinetuning(BaseFinetuning): def __init__(self, unfreeze_at_epoch5, layers_per_epoch2): super().__init__() self.unfreeze_at_epoch unfreeze_at_epoch self.layers_per_epoch layers_per_epoch def freeze_before_training(self, pl_module): # Freeze the entire backbone initially self.freeze(pl_module.backbone) def finetune_function(self, pl_module, epoch, optimizer): # Gradually unfreeze layers if epoch self.unfreeze_at_epoch: layers_to_unfreeze min( self.layers_per_epoch, len(list(pl_module.backbone.children())) ) # Unfreeze from the top layers down backbone_children list(pl_module.backbone.children()) for layer in backbone_children[-layers_to_unfreeze:]: self.unfreeze_and_add_param_group( layer, optimizer, lr1e-4 )这个从顶层往下逐层解冻的自定义策略正是BaseFinetuning开放性的体现。基类提供了一组可直接复用的静态工具方法均可在 finetuning.py 中查看实现方法作用flatten_modules(modules)把模块/模块迭代器展平为叶子模块与自身持有参数的父模块列表正确处理ModuleDict与嵌套结构freeze(modules, train_bnTrue)冻结指定模块参数train_bnTrue时 BN 层保持可训练make_trainable(modules)解冻模块参数并恢复 BN 的track_running_statsfilter_params(modules, train_bn, requires_grad)按requires_grad过滤生成参数迭代器train_bnFalse时跳过 BN 层filter_on_optimizer(optimizer, params)排除已存在于优化器其他 param_group 中的参数避免重复优化并给出警告unfreeze_and_add_param_group(modules, optimizer, lrNone, initial_denom_lr10.0, train_bnTrue)解冻模块并将参数作为新 param_group 加入优化器未指定lr时使用param_groups[0][lr] / initial_denom_lr两个使用要点来自源码与测试configure_optimizers中务必按requires_grad过滤参数。基类文档示例推荐Adam(filter(lambda p: p.requires_grad, self.parameters()))。若你把已冻结或将被回调解冻的参数同时放进了初始优化器filter_on_optimizer会检测到重复并发出UserWarningThe provided params to be frozen already exist within another group of this optimizer提示检查configure_optimizers的写法——该场景在 test_finetuning_callback_warning 中被验证。断点续训时参数组会被自动重建。BaseFinetuning.state_dict保存了_internal_optimizer_metadataparam_group 的参数名映射而非张量引用on_fit_start在优化器重建后通过_apply_mapping_to_param_groups把参数名映射回新模型的参数并恢复 param_groupsfinetuning.py。测试test_callbacks_restoretest_finetuning_callback.py断言了元数据中保存的params字段如[layer.3.weight, layer.3.bias]以及恢复后每个 param_group 的完整超参字典test_callbacks_restore_backbone则验证了带BackboneFinetuning的训练中断后从 checkpoint 继续训练的可行性。七、总结选择适合你的微调路线围绕迁移学习PyTorch Lightning 提供了从手写到全自动的完整梯度简单复用任何nn.Module都可以直接嵌入LightningModule预训练LightningModule通过load_from_checkpoint(...)加载后配合内置freeze()/unfreeze()控制参数是否可训练适合一次性设定冻结/解冻状态的场景。自动微调BackboneFinetuning回调实现冻结主干 → 指定 epoch 解冻 → 学习率分层调度 → 对齐 head 学习率的完整流程仅需约定模型含backbone与head两个属性适合大多数 CV/NLP 微调任务。完全自定义继承BaseFinetuning重写freeze_before_training与finetune_function两个钩子配合freeze、make_trainable、unfreeze_and_add_param_group等工具方法可精确控制每一层在每一轮的冻结/解冻与学习率适合复杂的分层训练策略。无论选择哪条路线所有微调状态包括动态新增的 optimizer param_group 元数据都会随 Lightning checkpoint 一并保存与恢复确保断点续训与生产部署的无缝衔接。【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000 GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考