MMPretrain项目中的训练调度策略详解

MMPretrain项目中的训练调度策略详解

mmpretrain OpenMMLab Pre-training Toolbox and Benchmark 项目地址: https://gitcode.com/gh_mirrors/mm/mmpretrain

前言

在深度学习模型训练过程中，优化策略和参数调度是影响模型性能的关键因素。MMPretrain作为一个强大的深度学习预训练框架，提供了灵活且全面的训练调度配置方案。本文将深入解析MMPretrain中的训练调度机制，帮助开发者根据实际需求定制优化策略。

优化器配置

基础优化器设置

MMPretrain支持PyTorch提供的所有优化器，通过optim_wrapper字段进行配置。这是一个典型的SGD优化器配置示例：

optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=dict(type='SGD', lr=0.0003, weight_decay=0.0001)
)

其中type指定优化器类型，lr设置学习率，weight_decay控制权重衰减强度。对于Adam优化器，可以这样配置：

optim_wrapper = dict(
    optimizer=dict(
        type='Adam',
        lr=0.001,
        betas=(0.9, 0.999),
        eps=1e-08,
        weight_decay=0,
        amsgrad=False)
)

混合精度训练

混合精度训练(AMP)可以显著减少显存占用并加速训练过程。启用方式非常简单：

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

或者在训练脚本中直接添加--amp参数。

参数级精细配置

不同类型参数差异化配置

MMPretrain提供了paramwise_cfg参数来实现对不同类型参数的差异化配置。例如，将归一化层的权重衰减设为0：

optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.8, weight_decay=1e-4),
    paramwise_cfg=dict(norm_decay_mult=0.))

支持的配置选项包括：

• bias_lr_mult：偏置项学习率乘数
• bias_decay_mult：偏置项权重衰减乘数
• norm_decay_mult：归一化层权重衰减乘数
• dwconv_decay_mult：深度可分离卷积权重衰减乘数

特定参数组配置

通过custom_keys可以针对特定参数组进行配置：

optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.01, weight_decay=0.0001),
    paramwise_cfg=dict(
        custom_keys={
            'backbone.layer0': dict(lr_mult=0, decay_mult=0),
            'backbone': dict(lr_mult=1),
            'head': dict(lr_mult=0.1)
        }))

梯度处理策略

梯度裁剪

梯度裁剪可防止梯度爆炸问题，配置方式如下：

optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.01, weight_decay=0.0001),
    clip_grad=dict(max_norm=35, norm_type=2))

梯度累积

当显存不足时，梯度累积是扩大等效batch size的有效方法：

train_dataloader = dict(batch_size=64)
optim_wrapper = dict(
    optimizer=dict(type='SGD', lr=0.01, weight_decay=0.0001),
    accumulative_counts=4)

这相当于使用256的batch size进行训练。

参数调度策略

学习率调度

MMPretrain支持多种学习率调度策略，包括MultiStepLR、CosineAnnealingLR等。

单调度器配置示例：

param_scheduler = dict(
    type='MultiStepLR',
    by_epoch=True,
    milestones=[100, 150],
    gamma=0.1)

多调度器组合（如warmup）：

param_scheduler = [
    dict(type='LinearLR',
        start_factor=0.001,
        by_epoch=False,
        end=50),
    dict(type='MultiStepLR',
        by_epoch=True,
        milestones=[8, 11],
        gamma=0.1)
]

动量调度

动量调度与学习率调度类似，可以配合使用：

param_scheduler = [
    dict(type='LinearLR', ...),
    dict(type='LinearMomentum',
         start_factor=0.001,
         by_epoch=False,
         begin=0,
         end=1000)
]

自定义优化策略

添加新优化器

1. 实现优化器类：

@OPTIMIZERS.register_module()
class MyOptimizer(Optimizer):
    def __init__(self, a, b, c):
        ...

2. 在配置中使用：

optim_wrapper = dict(
    optimizer=dict(type='MyOptimizer', a=a_value, b=b_value, c=c_value))

自定义优化器构造器

实现自定义构造器：

@OPTIM_WRAPPER_CONSTRUCTORS.register_module()
class MyOptimWrapperConstructor:
    def __init__(self, optim_wrapper_cfg, paramwise_cfg=None):
        ...

在配置中指定：

optim_wrapper = dict(
    constructor=dict(type='MyOptimWrapperConstructor'),
    optimizer=...,
    paramwise_cfg=...,
)

总结

MMPretrain提供了高度灵活的训练调度配置方案，从基础的优化器选择到复杂的参数级调度策略，开发者可以根据实际需求进行精细调整。通过合理配置这些参数，可以显著提升模型训练效果和收敛速度。

mmpretrain OpenMMLab Pre-training Toolbox and Benchmark 项目地址: https://gitcode.com/gh_mirrors/mm/mmpretrain