MMSelfSup项目自定义运行时配置详解

最新推荐文章于 2025-06-11 09:20:01 发布

符凡言Elvis

最新推荐文章于 2025-06-11 09:20:01 发布

阅读量361

点赞数 3

CC 4.0 BY-SA版权

本文链接：https://blog.youkuaiyun.com/gitblog_00526/article/details/148579071

MMSelfSup项目自定义运行时配置详解

mmselfsup OpenMMLab Self-Supervised Learning Toolbox and Benchmark 项目地址: https://gitcode.com/gh_mirrors/mm/mmselfsup

前言

在深度学习模型训练过程中，灵活配置训练流程是提高模型性能的关键。MMSelfSup作为自监督学习框架，提供了丰富的运行时配置选项，允许开发者根据需求自定义训练循环、钩子机制、优化器和学习率调度器等核心组件。本文将深入解析这些自定义配置的实现方法。

训练循环(Loop)配置

训练循环是控制模型训练流程的核心组件，MMSelfSup基于MMEngine的Loop机制提供了高度可定制的训练流程控制。

基本循环类型

MMSelfSup支持两种主要循环类型：

基于轮次的循环(EpochBasedTrainLoop)：这是最常用的循环类型，按照完整数据集遍历次数(epoch)进行训练

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

基于迭代次数的循环(IterBasedTrainLoop)：按照迭代次数进行训练，适用于大数据集场景

自定义循环实现

当内置循环无法满足需求时，可以继承BaseLoop类实现自定义循环。自定义循环需要实现以下关键方法：

run(): 定义循环的主要逻辑
train(): 训练过程的具体实现
val(): 验证过程的具体实现

钩子(Hook)机制详解

钩子机制是MMSelfSup实现训练过程扩展的核心设计，允许在不修改主流程代码的情况下插入自定义逻辑。

钩子类型与执行点

MMSelfSup提供了丰富的钩子执行点，覆盖训练全生命周期：

训练阶段：before_train, after_train, before_train_epoch, after_train_epoch, before_train_iter, after_train_iter
验证阶段：类似训练阶段的对应钩子点
测试阶段：类似训练阶段的对应钩子点

自定义钩子实现三步法

第一步：创建新钩子类

继承Hook基类并实现目标方法，例如实现一个动态调整超参数的钩子：

@HOOKS.register_module()
class DynamicParamHook(Hook):
    """动态调整模型超参数的钩子"""
    
    def __init__(self, factor_a: int, factor_b: int):
        self.factor_a = factor_a
        self.factor_b = factor_b
    
    def before_train_iter(self, runner, batch_idx, data_batch=None):
        current_iter = runner.iter
        model = get_model(runner.model)
        model.hyper_param = self.factor_a * current_iter + self.factor_b

第二步：注册新钩子

在对应模块的__init__.py中导入并注册新钩子：

from .dynamic_param_hook import DynamicParamHook

__all__ = [..., 'DynamicParamHook']

第三步：配置使用钩子

在配置文件中添加自定义钩子：

custom_hooks = [
    dict(type='DynamicParamHook', factor_a=0.1, factor_b=0.5)
]

钩子优先级控制

MMSelfSup支持通过priority参数控制钩子执行顺序：

HIGHEST: 最高优先级
VERY_HIGH: 非常高
HIGH: 高
ABOVE_NORMAL: 高于正常
NORMAL (默认): 正常
BELOW_NORMAL: 低于正常
LOW: 低
VERY_LOW: 非常低
LOWEST: 最低

优化器配置进阶

优化器配置是模型训练的关键环节，MMSelfSup支持完整的PyTorch优化器生态。

基础优化器配置

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

优化器包装器

优化器包装器提供统一接口支持不同训练模式：

标准包装器(OptimWrapper):

optim_wrapper = dict(type='OptimWrapper', optimizer=optimizer)

混合精度包装器(AmpOptimWrapper):

optim_wrapper = dict(type='AmpOptimWrapper', optimizer=optimizer)

参数级优化配置

通过paramwise_cfg可以实现细粒度参数优化控制：

optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=optimizer,
    paramwise_cfg=dict(
        custom_keys={
            'backbone': dict(lr_mult=0.1),  # 骨干网络学习率乘数
            'norm': dict(decay_mult=0),     # 归一化层权重衰减
            'bias': dict(lr_mult=2)         # 偏置项学习率乘数
        }))

典型配置示例

MAE优化配置:

optimizer = dict(type='AdamW', lr=1.5e-4, betas=(0.9, 0.95), weight_decay=0.05)
optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=optimizer,
    paramwise_cfg=dict(
        custom_keys={
            'ln': dict(decay_mult=0.0),       # LayerNorm无权重衰减
            'pos_embed': dict(decay_mult=0.),  # 位置编码无权重衰减
            'cls_token': dict(decay_mult=0.)   # 分类token无权重衰减
        }))

SimCLR优化配置:

optimizer = dict(type='LARS', lr=0.3, momentum=0.9, weight_decay=1e-6)
optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=optimizer,
    paramwise_cfg=dict(
        custom_keys={
            'bn': dict(lars_exclude=True),  # BN层不使用LARS优化
            'bias': dict(lars_exclude=True) # 偏置项不使用LARS优化
        }))

学习率调度器配置

学习率调度是训练过程中的关键超参数，MMSelfSup支持丰富多样的调度策略。

组合调度策略

param_scheduler = [
    # 初始阶段线性warmup
    dict(
        type='LinearLR',
        start_factor=1e-4,
        by_epoch=True,
        begin=0,
        end=10,
        convert_to_iter_based=True),
    # 主训练阶段余弦退火
    dict(
        type='CosineAnnealingLR',
        T_max=190,  # T_max = max_epochs - warmup_epochs
        by_epoch=True,
        begin=10,
        end=200)
]