MMDetection项目中的实用Hook机制详解-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00199/article/details/148359896

MMDetection项目中的实用Hook机制详解

mmdetection open-mmlab/mmdetection: 是一个基于 PyTorch 的人工智能物体检测库，支持多种物体检测算法和工具。该项目提供了一个简单易用的人工智能物体检测库，可以方便地实现物体的检测和识别，同时支持多种物体检测算法和工具。项目地址: https://gitcode.com/gh_mirrors/mm/mmdetection

什么是Hook机制

Hook（钩子）机制是深度学习中一种强大的扩展设计模式，它允许开发者在训练流程的关键节点插入自定义操作，而无需修改核心代码。在MMDetection目标检测框架中，Hook机制为开发者提供了灵活干预训练过程的能力。

MMDetection中的内置Hook

MMDetection提供了多种实用的内置Hook，下面介绍几个核心Hook的功能和使用方法：

1. 损失有效性检查Hook (CheckInvalidLossHook)

这个Hook会在训练过程中定期检查损失值是否变为NaN或无限大。当检测到异常损失时，会立即终止训练并提示错误信息。

典型应用场景：

训练初期出现梯度爆炸
学习率设置过高导致数值不稳定
数据中存在异常样本

2. 类别数量验证Hook (NumClassCheckHook)

该Hook用于验证模型配置中的类别数量是否与数据集的实际类别数量一致，避免常见的配置错误。

重要性：

防止因类别数不匹配导致的模型维度错误
确保检测头输出层与数据集匹配
在训练开始前进行早期验证

3. 内存分析Hook (MemoryProfilerHook)

内存分析Hook是开发者调试内存问题的利器，它可以记录以下关键内存信息：

系统虚拟内存使用情况
交换空间(swap)使用情况
当前训练进程的内存占用

使用前提：需要先安装依赖包：

pip install memory_profiler psutil

配置示例：

custom_hooks = [
    dict(type='MemoryProfilerHook', interval=50)  # 每50次迭代记录一次
]

输出示例：

系统共有250GB内存(246360MB + 9407MB)和8GB交换内存(5740MB + 2452MB)
当前内存使用率4.4%(9407MB)，交换内存使用率29.9%(5740MB)
当前训练进程占用5434MB内存

4. 同步归一化Hook (SyncNormHook)

在多GPU训练场景下，该Hook确保BatchNorm层的统计信息在所有GPU间同步，提高训练稳定性。

5. YOLOX专用Hook

针对YOLOX系列模型，MMDetection提供了两个专用Hook：

YOLOXLrUpdaterHook：YOLOX特有的学习率调整策略
YOLOXModeSwitchHook：处理YOLOX训练过程中的模式切换

如何自定义Hook

MMDetection的Hook系统提供了20个可插入点，覆盖训练全流程：

主要插入点分类

全局生命周期点：
- before_run：训练开始前
- after_run：训练结束后
训练过程点：
- before_train/after_train：训练阶段前后
- before_train_epoch/after_train_epoch：每个epoch前后
- before_train_iter/after_train_iter：每次迭代前后
验证过程点：
- 类似训练过程，包含val相关的各个阶段
测试过程点：
- 包含test相关的各个阶段
检查点相关点：
- before_save_checkpoint/after_save_checkpoint：保存模型前后

自定义Hook实现步骤

以创建一个损失检查Hook为例：

继承基础Hook类并实现关键方法
使用装饰器注册Hook
在配置中启用Hook

示例代码：

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

@HOOKS.register_module()
class CustomLossCheckHook(Hook):
    def __init__(self, interval=50):
        self.interval = interval
        
    def after_train_iter(self, runner, batch_idx, data_batch, outputs):
        if self.every_n_train_iters(runner, self.interval):
            if not torch.isfinite(outputs['loss']):
                runner.logger.warning('发现非正常损失值!')
                # 可添加自定义处理逻辑