Hi-FT/ERD项目自定义训练运行时配置详解

前言

在深度学习模型训练过程中，合理配置训练运行时参数对模型性能有着至关重要的影响。本文将深入讲解如何在Hi-FT/ERD项目中自定义训练配置，包括优化器设置、训练策略调整、训练循环定制以及钩子功能扩展等高级技巧。

优化器配置详解

基础优化器配置

Hi-FT/ERD项目采用OptimWrapper统一管理优化相关配置，主要包含三个核心部分：

1. 优化器(optimizer)：定义基础优化算法及参数
2. 参数级配置(paramwise_cfg)：实现不同参数组的差异化设置
3. 梯度裁剪(clip_grad)：控制梯度更新幅度

典型配置示例如下：

optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=dict(
        type='AdamW',  # 使用AdamW优化器
        lr=0.0001,     # 基础学习率
        weight_decay=0.05,  # 权重衰减系数
        eps=1e-8,      # 数值稳定项
        betas=(0.9, 0.999)),  # 动量参数
    
    paramwise_cfg=dict(
        custom_keys={
            'backbone': dict(lr_mult=0.1, decay_mult=1.0),  # 主干网络学习率缩小10倍
        },
        norm_decay_mult=0.0),  # 归一化层权重衰减系数
    
    clip_grad=dict(max_norm=0.01, norm_type=2)  # 梯度裁剪阈值
)

内置优化器使用

项目支持PyTorch所有原生优化器，只需简单修改配置即可切换。例如使用SGD优化器：

optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001)
)

自定义优化器实现

当内置优化器无法满足需求时，可以扩展自定义优化器：

1. 定义优化器类：继承torch.optim.Optimizer并实现核心方法

from mmdet.registry import OPTIMIZERS
from torch.optim import Optimizer

@OPTIMIZERS.register_module()
class MyOptimizer(Optimizer):
    def __init__(self, params, a, b, c, **kwargs):
        # 实现初始化逻辑
        pass
    
    def step(self, closure=None):
        # 实现参数更新逻辑
        pass

2. 注册优化器：确保模块被正确导入

   • 方式一：修改__init__.py文件显式导入
   • 方式二：配置中通过custom_imports动态导入

3. 配置使用：在optimizer域指定自定义优化器

optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=dict(type='MyOptimizer', a=1.0, b=0.5, c=0.1)
)

高级优化技巧

1. 梯度裁剪：稳定训练过程

optim_wrapper = dict(clip_grad=dict(max_norm=35, norm_type=2))

2. 动量调度：配合学习率调度加速收敛

param_scheduler = [
    # 学习率调度
    dict(type='CosineAnnealingLR', T_max=8, eta_min=lr*10),
    # 动量调度  
    dict(type='CosineAnnealingMomentum', T_max=8, eta_min=0.85)
]

训练策略定制

学习率调度策略

项目支持多种学习率调整策略：

1. 多项式衰减策略：

param_scheduler = [
    dict(type='PolyLR', power=0.9, eta_min=1e-4, begin=0, end=8)
]

2. 余弦退火策略：

param_scheduler = [
    dict(type='CosineAnnealingLR', T_max=8, eta_min=lr*1e-5)
]

训练循环配置

支持两种基础训练循环模式：

1. 基于轮次的循环(EpochBasedTrainLoop)：

train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=12)

2. 基于迭代的循环(IterBasedTrainLoop)：

train_cfg = dict(
    type='IterBasedTrainLoop',
    max_iters=368750,
    val_interval=5000,
    dynamic_intervals=[(365001, 368750)]  # 动态调整验证间隔
)

钩子功能扩展

自定义钩子实现

1. 定义钩子类：继承Hook基类并实现所需方法

from mmengine.hooks import Hook
from mmdet.registry import HOOKS

@HOOKS.register_module()
class MyHook(Hook):
    def __init__(self, a, b):
        pass
    
    def before_train_epoch(self, runner):
        # 训练epoch前执行逻辑
        pass
    
    def after_train_iter(self, runner, batch_idx, data_batch, outputs):
        # 训练迭代后执行逻辑
        pass

2. 注册与使用：

custom_hooks = [
    dict(type='MyHook', a=1, b=2, priority='NORMAL')
]

内置钩子配置

1. 模型检查点钩子：

default_hooks = dict(
    checkpoint=dict(
        type='CheckpointHook',
        interval=1,
        max_keep_ckpts=3,  # 最多保留3个检查点
        save_optimizer=True  # 保存优化器状态
    )
)

2. 日志记录钩子：

default_hooks = dict(
    logger=dict(type='LoggerHook', interval=50)  # 每50次迭代记录一次
)

3. 可视化钩子：

default_hooks = dict(
    visualization=dict(type='DetVisualizationHook', draw=True)
)

vis_backends = [
    dict(type='LocalVisBackend'),  # 本地可视化
    dict(type='TensorboardVisBackend')  # TensorBoard支持
]

结语

通过本文的详细介绍，相信您已经掌握了在Hi-FT/ERD项目中自定义训练运行时的各项高级技巧。合理配置优化策略、训练循环和监控钩子，可以显著提升模型训练效率和最终性能。建议根据具体任务需求，灵活组合文中介绍的各种配置方法，找到最适合您项目的训练方案。