SuperGradients 学习率分配实战：initial_lr 映射、finetune 冻结与 YoloNAS 微调指南-平芜编程栈

SuperGradients 学习率分配实战：initial_lr 映射、finetune 冻结与 YoloNAS 微调指南

【免费下载链接】super-gradientsEasily train or fine-tune SOTA computer vision models with one open source training library. The home of Yolo-NAS.项目地址: https://gitcode.com/GitHub_Trending/su/super-gradients

导读

在迁移学习与微调场景中，给神经网络的不同层（或不同参数组）分配差异化的初始学习率，是兼顾"保留预训练特征"与"快速适配新任务"的关键手段。SuperGradients（SG）通过training_params中的initial_lr与finetune两个超参数，提供了一套从"全模型统一学习率"到"按层前缀精准分配"再到"一行代码自动冻结骨干网络"的完整方案。读完本文，你将掌握如何在 SuperGradients 中使用标量或映射形式的initial_lr、如何通过finetune=True对 YoloNAS 等检测模型做"只训练检测头"的微调，并理解其底层在 optimizer_utils.py 中的参数分组实现。

一、理解`initial_lr`：全局标量 vs. 前缀映射

在 SG 中，initial_lr是训练超参数（training hyperparameter），决定优化器创建时各参数组的初始学习率。它支持两种形态：

形态	写法	效果
标量（float）	`"initial_lr": 0.01`	所有参数使用统一学习率
映射（dict）	`"initial_lr": {"conv": 0.001, "default": 0.}`	按参数名前缀分组，各组使用不同学习率

从源码看，这一分支逻辑位于 initialize_param_groups：当lr是float或int时，模型全部命名参数被放进单个"default"分组；否则调用 separate_lr_groups 按前缀切分。随后在 build_optimizer 中，这些分组会被转换为 torch 优化器（SGD/Adam/RMSprop 等）的param_groups。

提示：initial_lr也接受Mapping形式与浮点数之外的 int 数值，如initialize_param_groups中对isinstance(lr, float) or isinstance(lr, int)的判断所示。

二、使用标量`initial_lr`：全模型统一学习率

当initial_lr是单个浮点数时，它对模型所有参数一视同仁。适合从头训练或整体微调预训练模型：

from super_gradients import Trainer # Define training parameters training_params = { "initial_lr": 0.01, "loss": "cross_entropy", # ... other training parameters } # Initialize the Trainer trainer = Trainer("simple_net_training") # Define model model = ... # Define data loaders train_dataloader = ... test_dataloader = ... # Train the model trainer.train(model, training_params, train_dataloader, test_dataloader)

在 sg_trainer.py 中，Trainer 通过build_optimizer(net=unwrap_model(self.net), lr=self.training_params.initial_lr, training_params=self.training_params)创建优化器。若你在training_params中直接传入实例化的 optimizer，则要求initial_lr为None，否则会抛出RuntimeError（见 sg_trainer.py）。

另外，如果使用外部实例化的优化器，SG 还会通过 get_initial_lr_from_optimizer 反向提取其param_groups中的学习率来填充initial_lr，并记录到日志中。

三、使用映射`initial_lr`：按层前缀精准分配

当initial_lr是字典时，SG 允许你对不同层使用不同学习率，适用于微调预训练模型或对网络不同部件采用差异化学习率。核心规则如下：

前缀匹配：字典的每个 key 作为模型命名参数（named_parameters）名称的前缀，所有参数名以该前缀开头的参数归入该分组，并使用对应的学习率。匹配通过name.startswith(group_name)实现（见 separate_lr_groups）。
defaultkey 是必需的：它为不匹配任何前缀的参数提供兜底学习率。源码中若lr_dict.get("default", None) is None会直接抛出RuntimeError("When passing initial_lr as dictionary, must pass 'default'.")（见 separate_lr_groups）。
冻结参数：给某个前缀分配学习率0即可冻结该部分参数。separate_lr_groups会将lr == 0的分组参数逐个设置param.requires_grad = False（见 separate_lr_groups），这些参数不会进入优化器参数组，训练时不会被更新。
default为 0 时同理冻结：所有未匹配前缀的参数也会被冻结（见 separate_lr_groups）。

下面示例中，conv1与conv2会被冻结（学习率 0），fc1与fc2以 0.001 的初始学习率训练：

import torch.nn as nn from super_gradients import Trainer class SimpleNet(nn.Module): def __init__(self): super(SimpleNet, self).__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 50, 5) self.fc1 = nn.Linear(50 * 4 * 4, 500) self.fc2 = nn.Linear(500, 10) def forward(self, x): x = nn.functional.relu(self.conv1(x)) x = nn.functional.relu(self.conv2(x)) x = x.view(-1, 50 * 4 * 4) x = nn.functional.relu(self.fc1(x)) x = self.fc2(x) return x trainer = Trainer("simple_net_training") # Define model model = SimpleNet() # Define data loaders train_dataloader = ... test_dataloader = ... # Define training parameters training_params = { "initial_lr": {"conv": 0.001, "default": 0.}, "loss": "cross_entropy", # ... other training parameters } # Train the model trainer.train(model, training_params, train_dataloader, test_dataloader)

注意："conv"前缀会同时匹配conv1和conv2；若只想匹配某个具体层，使用更长、更具体的前缀（如"fc1"）即可。由于前缀匹配采用startswith，设计 key 时应避免出现意外的交叉匹配。单元测试 test_lr_assignment.py 验证了分组必须"覆盖全部参数"且"分组之间无交集"这两个约束。

四、`finetune`特性：一行代码冻结骨干网络

对于不熟悉网络各部件命名的用户，SG 提供了finetune参数：将其设为True后，模型会根据自身实现的get_finetune_lr_dict方法自动冻结一部分网络，只训练另一部分。该特性对 SG model zoo 中所有实现了get_finetune_lr_dict方法的模型均可用。

例如，下面的代码会让 YoloNAS 的检测头（heads）以 0.01 的初始学习率训练，而网络其余部分被冻结：

from super_gradients import Trainer, models from super_gradients.common.object_names import Models trainer = Trainer("simple_net_training") # Define model model = models.get(Models.YOLO_NAS_S, pretrained_weights="coco", num_classes=2) # Define data loaders train_dataloader = ... test_dataloader = ... # Define training parameters training_params = { "initial_lr": 0.01, "finetune": True, # ... other training parameters } # Train the model trainer.train(model, training_params, train_dataloader, test_dataloader)

4.1`finetune`的工作原理

finetune=True的执行链路清晰且可验证：

接口约定：模型需实现SupportsFineTune接口，其唯一方法get_finetune_lr_dict(self, lr: float) -> Dict[str, float]返回一个与initial_lr映射同构的字典，将学习率映射到网络中"不被冻结的部分"（见 module_interfaces.py）。
运行时替换：在 build_optimizer 中，当training_params.finetune为真时：若模型不是SupportsFineTune实例，则打印警告说明finetune=True不生效；若模型实现了该接口但initial_lr不是标量，则抛出RuntimeError("When training with fine_tune=True, initial_lr must be a scalar.")；随后用net.get_finetune_lr_dict(lr)的返回值覆盖lr，并同步写回training_params.initial_lr，同时打印日志Training with finetune=True: setting initial_lr to predefined mapping ...。
YoloNAS 的实现：在 customizable_detector.py 中，YoloNAS 的get_finetune_lr_dict返回：

def get_finetune_lr_dict(self, lr: float): return {"heads": lr, "default": 0}

即：只有名字以heads开头的参数（检测头）使用传入的lr训练，其余全部冻结。

4.2 约束与注意事项

initial_lr必须是标量：finetune设计上只与"未设置或为浮点数"的initial_lr配合使用；若initial_lr已是映射（dict），启用finetune会报错。
模型必须实现接口：如果模型未实现get_finetune_lr_dict，finetune=True只会产生告警而不会生效（见 build_optimizer 的warnings.warn分支）。
与优化器参数组的交互：finetune生成的映射最终同样进入separate_lr_groups流程，因此 "default": 0 会冻结非 heads 参数；若同时开启zero_weight_decay_on_bias_and_bn，冻结参数也不会进入 weight decay 分组（见 separate_zero_wd_params_groups_for_optimizer）。

五、从测试与配置看最佳实践

5.1 测试用例验证冻结效果

test_finetune.py 用真实训练断言了finetune的冻结行为，例如 YoloNAS 用例（见 test_finetune.py）：

check_models_have_same_weights(net_before_train.backbone, net.backbone, skip_bn_stats=True)为真：训练后 backbone 权重不变（已冻结）；
check_models_have_same_weights(net_before_train.neck, net.neck, skip_bn_stats=True)为真：neck 权重不变；
check_models_have_same_weights(net_before_train.heads, net.heads)为假：heads 权重被更新。

同样的模式也覆盖了 PP-YoloE、YoloX、DDRNet、PP-LiteSeg、RegSeg、SegFormer、STDC、BEiT、EfficientNet、MobileNet、RegNet、RepVGG、ResNet 等模型（见 test_finetune.py 各test_train_with_finetune_*用例），可据此判断你使用的模型是否支持finetune以及冻结范围。

5.2 与调度器、权重衰减等超参数的配合

initial_lr只是"初始"学习率，后续的衰减由lr_mode（如CosineLRScheduler、StepLRScheduler、PolyLRScheduler、FunctionLRScheduler等）接管。默认训练参数定义在 default_train_params.yaml：

optimizer: SGD # 优化算法，可选 'Adam'、'SGD'、'RMSProp' optimizer_params: {} # 优化器初始化参数，如 {"weight_decay": 0.0001, "momentum": 0.9} zero_weight_decay_on_bias_and_bn: False # 是否对 BN 与 bias 关闭 weight decay finetune: False # 是否冻结模型固定部分，仅对实现 get_finetune_lr_dict 的模型有效

在该配置文件中，finetune的注释与本文一致：finetune为 True 时冻结模型固定部分，模型类方法get_finetune_lr_dict返回"将 lr 映射到网络未被冻结部分"的字典，用法与initial_lr映射一致。

几个实操建议：

微调检测模型（如 YoloNAS、YoloX、PP-YoloE）：设置finetune=True并配合较小的initial_lr（如5e-4，参考 test_finetune.py），配合lr_mode: cosine与 AdamW。
微调分割模型（如 DDRNet、PP-LiteSeg、SegFormer、STDC、RegSeg）：finetune=True会冻结除分割头之外的网络，测试中通常配合PolyLRScheduler与 SGD（momentum=0.9、weight_decay=5e-4，见 test_finetune.py）。
分类模型微调（如 ResNet、EfficientNet、RegNet、MobileNet、BEiT、RepVGG）：finetune冻结除分类头（linear/classifier/head/_fc）之外的部分，测试中常用StepLRScheduler与较大的initial_lr（如 0.6），并通过lr_updates、lr_decay_factor控制衰减（见 test_finetune.py）。
需要精确控制层组学习率：使用initial_lr映射，务必包含"default"key；需要冻结时令目标前缀的 lr 为 0。

六、总结

SuperGradients 的学习率分配机制可以总结为一条清晰的主线：

标量initial_lr：全体参数统一学习率，最简用法；
映射initial_lr：按named_parameters前缀分组，实现差异化学习率与参数冻结，default兜底；
finetune=True：调用模型实现的get_finetune_lr_dict，自动生成映射（如 YoloNAS 的{"heads": lr, "default": 0}），实现"只训练检测头/分割头/分类头"的即插即用微调。

三者最终都汇入 initialize_param_groups → separate_lr_groups →build_optimizer的参数分组管线，并被转换为 torch 优化器的param_groups。无论你是在 SuperGradients 中微调 YoloNAS 检测模型、DDRNet 系列分割模型，还是 ResNet 系列分类模型，都可以按本文提供的三种模式灵活控制训练过程中每一组参数的学习率与冻结状态。

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

SuperGradients 学习率分配实战：initial_lr 映射、finetune 冻结与 YoloNAS 微调指南