MMsegmentation项目自定义运行配置指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00126/article/details/148419449

MMsegmentation项目自定义运行配置指南

mmsegmentation OpenMMLab Semantic Segmentation Toolbox and Benchmark. 项目地址: https://gitcode.com/gh_mirrors/mm/mmsegmentation

前言

在深度学习模型训练过程中，灵活地定制训练流程和优化策略是提升模型性能的关键。本文将详细介绍如何在MMsegmentation项目中实现自定义的运行配置，包括钩子(Hook)、优化器(Optimizer)和优化器封装构造器(OptimWrapperConstructor)的定制方法。

自定义训练钩子(Hook)

钩子的概念与作用

钩子是深度学习训练过程中的回调机制，允许开发者在训练的不同阶段（如每次迭代前后、每个epoch前后等）插入自定义逻辑。MMsegmentation基于MMEngine提供了丰富的内置钩子，但当我们需要实现特殊训练逻辑时，就需要自定义钩子。

实现自定义钩子的步骤

1. 创建新的钩子类

下面是一个实现动态调整模型超参数的钩子示例：

from mmengine.hooks import Hook
from mmengine.model import is_model_wrapper
from mmseg.registry import HOOKS

@HOOKS.register_module()
class DynamicParamHook(Hook):
    """动态调整模型超参数的钩子
    
    该钩子会根据训练迭代次数线性调整模型的hyper_parameter参数
    """
    
    def __init__(self, multiplier: int, base_value: int) -> None:
        self.multiplier = multiplier  # 线性变化的系数
        self.base_value = base_value  # 基础值
        
    def before_train_iter(self, runner, batch_idx: int, data_batch=None) -> None:
        current_iter = runner.iter
        # 处理模型被封装的情况
        model = runner.model.module if is_model_wrapper(runner.model) else runner.model
        # 线性调整超参数
        model.hyper_parameter = self.multiplier * current_iter + self.base_value

2. 导入钩子模块

有两种方式使系统能够识别自定义钩子：

修改__init__.py文件：在mmseg/engine/hooks/__init__.py中添加导入语句
使用配置文件导入：在配置文件中添加custom_imports配置项

3. 配置使用钩子

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

custom_hooks = [
    dict(
        type='DynamicParamHook',
        multiplier=0.01,  # 每次迭代增加0.01
        base_value=0.1,    # 初始值为0.1
        priority='ABOVE_NORMAL'  # 执行优先级
    )
]

自定义优化器实现

优化器定制场景

当内置优化器无法满足特定需求时，例如需要实现特殊的参数更新规则或支持新的优化算法时，就需要自定义优化器。

实现自定义优化器的步骤

1. 创建优化器类

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

@OPTIMIZERS.register_module()
class CustomAdam(Optimizer):
    """实现自定义Adam优化器变体
    
    参数:
        params (iterable): 可迭代的参数或参数组
        lr (float): 学习率
        beta1 (float): 梯度一阶矩估计的指数衰减率
        beta2 (float): 梯度二阶矩估计的指数衰减率
        epsilon (float): 数值稳定项
    """
    
    def __init__(self, params, lr=1e-3, beta1=0.9, beta2=0.999, epsilon=1e-8):
        defaults = dict(lr=lr, beta1=beta1, beta2=beta2, epsilon=epsilon)
        super().__init__(params, defaults)
        
    def step(self, closure=None):
        # 实现参数更新逻辑
        ...

2. 导入优化器模块

同样有两种导入方式：

修改mmseg/engine/optimizers/__init__.py文件
在配置文件中使用custom_imports

3. 配置使用优化器

optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=dict(
        type='CustomAdam',
        lr=0.001,
        beta1=0.9,
        beta2=0.999,
        epsilon=1e-8
    ),
    clip_grad=None
)

自定义优化器封装构造器

构造器的作用

优化器封装构造器用于更精细地控制不同网络层的优化策略，例如：

为不同层设置不同的学习率
对特定层（如BatchNorm）禁用权重衰减
实现分层学习率衰减策略

实现自定义构造器的步骤

1. 创建构造器类

from mmengine.optim import DefaultOptimWrapperConstructor
from mmseg.registry import OPTIM_WRAPPER_CONSTRUCTORS

@OPTIM_WRAPPER_CONSTRUCTORS.register_module()
class LayerwiseLRConstructor(DefaultOptimWrapperConstructor):
    """实现分层学习率衰减的优化器构造器
    
    参数:
        decay_rate (float): 学习率衰减率
        decay_layers (list): 需要衰减的层名称列表
    """
    
    def __init__(self, optim_wrapper_cfg, paramwise_cfg=None):
        super().__init__(optim_wrapper_cfg, paramwise_cfg)
        self.decay_rate = paramwise_cfg.get('decay_rate', 0.9)
        self.decay_layers = paramwise_cfg.get('decay_layers', [])
        
    def __call__(self, model):
        # 实现分层学习率设置逻辑
        ...

2. 导入构造器模块

同样可以通过修改__init__.py或使用custom_imports实现

3. 配置使用构造器

optim_wrapper = dict(
    type='OptimWrapper',
    constructor=dict(
        type='LayerwiseLRConstructor',
        paramwise_cfg=dict(
            decay_rate=0.95,
            decay_layers=['backbone']
        )
    ),
    optimizer=dict(
        type='AdamW',
        lr=0.0001,
        weight_decay=0.01
    )
)