Skip to main content

如何搭配MMDetection使用Weights & Biases

在本报告中,我们将使用MMDetection来训练一个目标检测模型,同时了解如何使用MMDetWandbHook。
Created on June 17|Last edited on June 17
本报告是作者Ayuth Thakur所写的"How to Use Weights & Biases with MMDetection"的翻译




目录



简介

我们今天的目标很简单:使用MMDetection训练一个目标检测模型并了解Weights & Biases能如何帮助您记录并验证指标、进行模型预测可视化、控制原始验证数据集版本等等。更确切地说,我们会在小型气球数据集上训练一个Mask RCNN模型。在详细介绍之前,我们先来聊聊MMDetection。
MMDetection是一款基于OpenMMLabs推出的PyTorch的流行开源资源库,用于执行目标检测任务。它提供了可进行组合的模块化API设计,便于您轻松构建自定义目标检测计划安排。此外还有超过100个预训练模型以及开箱即用的标准数据集。同时还支持单显卡和多显卡训练。您可在这里查看GitHub库。
W&B方面,本文将展示一个专门用于MMDetection的 MMDetWandbHook,可用于:
✅ 记录训练和评估指标。
✅ 记录不同版本模型检查点。
✅ 使用基本事实边界框记录多版本验证数据集。
✅ 对模型预测进行记录和可视化。
我们正式开始吧。

安装


这里我们将快速设置在自定义数据集上使用MMDetection进行训练所需使用的一切选项。如果您想使用可执行代码跟学,也可查看上方的Colab链接。

安装MMDetection

MMDetection对于MMCV库十分依赖。如果您对此不熟悉,MMCV是OpenMMLabs 旗下大多数项目(库)所使用的基础库。请注意,我们必须安装与给定 PyTorch版本兼容的正确版本的 MMCV。MMCV与MMDetection均可使用pip进行安装。可查看安装文档页面以了解更多详细信息。
安装所需的代码如下:
# Install PyTorch.
!pip install -qq torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html

# Install mmcv-full.
!pip install -qq mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.9.0/index.html

# Install mmdetection
!git clone https://github.com/open-mmlab/mmdetection
%cd mmdetection
!pip install -e .

安装Weights & Biases

您可使用适用于Python的pip 安装程序包轻松安装Weights & Biases。如果您还没有创建W&B账户,则可访问其注册页面创建一个。您可使用私密身份认证代币进行登录并认证一台机器,之后即可开始向W&B中记录运行状况(指标、数据等)。
!pip install -q --upgrade wandb

import wandb
wandb.login()

数据集


如上所述,这里我们会使用一个小型气球数据集。您可使用下方的代码片段下载该数据集:
!wget https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip
!unzip -q balloon_dataset.zip
为了支持新的数据格式,我们建议将标注转换为COCO格式或PASCAL VOC格式。如果您需要将标注转换为COCO格式,请离线进行并使用 CocoDataset 类。如果您需要将其转换为PASCAL格式,请使用 VOCDataset 类。这些类可使用配置系统来进行分配,下文有提及。您可查看与本报告相关的Colab笔记本,了解我们是如何转换气球数据集格式的。这里有更多关于自定义数据集使用方面的信息。
MMDetWandbHook 将自动记录验证数据,样本数量可通过使用 num_eval_samples 来确定。我们将在稍后深入研究 MMDetWandbHook,不过这里会提供有关数据集的简要预览。以下显示的面板是用于记录验证数据的W&B表格。



模型


MMDetection通过Model Zoo 提供了百余个预训练的目标检测器。这里是其文档页面
您可以轻松自定义模型的 backbone、neck、head、ROI以及loss,还可在这里了解有关自定义模型的更多信息。在本教程中,我们将使用预训练的模型检查点,并在自定义数据集上对其进行微调。现在我们来下载一下 checkpoints 目录中的模型:
!mkdir checkpoints
!wget -c https://download.openmmlab.com/mmdetection/v2.0/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth \
-O checkpoints/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth

配置


MMDetection严重依赖于配置系统。配置系统能够将 MMDetection的各个组件整合在一起��所有目标检测框架,比如 Faster R-CNN、Mask R-CNN、Cascade R-CNN、RPN、SSD、YOLO等等,都有其自己的配置文件。我们可以根据需要加载文件并修改一些方法。您可在此处进一步了解有关MMDetection配置系统的信息。
目前,大多数配置文件都会使用 config/__base__ 下的四个基本组件类型:datasetmodelscheduledefault_runtime。您可以直接修改文件中的方法,也可以加载配置文件并通过使用下方的 cfg 对象来使用必要方法:
config_file = 'mmdetection/configs/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_1x_coco.py'
cfg = Config.fromfile(config_file)
注意这里的 COCODataset 使用方法:

# Define type and path to the images.
cfg.dataset_type = 'COCODataset'

cfg.data.test.ann_file = 'balloon/val/annotation_coco.json'
cfg.data.test.img_prefix = 'balloon/val/'
cfg.data.test.classes = ('balloon',)

cfg.data.train.ann_file = 'balloon/train/annotation_coco.json'
cfg.data.train.img_prefix = 'balloon/train/'
cfg.data.train.classes = ('balloon',)

cfg.data.val.ann_file = 'balloon/val/annotation_coco.json'
cfg.data.val.img_prefix = 'balloon/val/'
cfg.data.val.classes = ('balloon',)
由于我们的数据集中只有一个类(当然是气球),我们将相应地指定模型来构建边界框 head 和 mask head。另外请注意我们在上一节中下载的模型检查点的使用方法:
# modify num classes of the model in box head and mask head
cfg.model.roi_head.bbox_head.num_classes = 1
cfg.model.roi_head.mask_head.num_classes = 1

# Use the pretrained model.
cfg.load_from = 'checkpoints/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth'
下面我们来定义训练配置。注意 checkpoint_config.intervalevaluation.interval 的使用方法,它们会在稍后派上大用场!
# The original learning rate (LR) is set for 8-GPU training.
# We divide it by 8 since we only use one GPU.
cfg.optimizer.lr = 0.02 / 8
cfg.lr_config.warmup = None
cfg.log_config.interval = 10

# Epochs
cfg.runner.max_epochs = 12

# Set seed thus the results are more reproducible
cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)

# ⭐️ Set the checkpoint interval.
cfg.checkpoint_config.interval = 4

# Set up working dir to save files and logs.
cfg.work_dir = './tutorial_exps'
cfg.device = 'cuda'

# ⭐️ Set the evaluation interval.
cfg.evaluation.interval = 2
为了使机器学习实验具有可复现性,我们必须对配置进行适当地追踪。我们将在后面的章节中看到如何使用 MMDetWandbHook 来将配置记录到W&B中。但在讨论 MMDetWandbHook 之前,我们先来简要看看MMDetection情景中的 Hook。

钩子

一般来说,“钩子”是在特定事件之后自动执行的函数。MMDetection提供各种标准“钩子”,可通过配置系统的 default_runtime 组件进行访问。下面列出了一些配置方法(对于本教程很重要):一般来说,“钩子”是在特定事件之后自动执行的函数。MMDetection提供各种标准“钩子”,可通过配置系统的 default_runtime 组件进行访问。下面列出了一些配置方法(对于本教程很重要):
  • log_config - 用于初始化多种记录器钩子,如:TextLoggerHookMMDetWandbHook 等等。
  • checkpoint_config - 用于初始化MMCV的 CheckpointHook。此钩子会按照上述配置部分中所定义的时间间隔来保存模型检查点。
  • evaluation - 用于根据可用显卡数量初始化MMCV的 EvalHookDistEvalHook。此钩子会按照上述配置部分中所定义的时间间隔来对模型进行验证数据评估。
通常而言,钩子都会有与其关联的优先级状态。在本例中,log_config 的优先级是 VERY_LOW,而其他钩子的则是 NORMAL

Weights & Biases的专用钩子


MMDetection现在提供了专门的 Weights和Biases钩子,名为 MMDetWandHook。借助这一专属钩子,您可以:
  • 记录训练和评估指标以及系统(处理器/显卡)指标
  • 将模型检查点保存为W&B artifact
  • 以交互式W&B表格的形式对验证数据集进行可视化
  • 以交互式W&B表格的形式对模型预测进行可视化
鉴于 MMDetWandbHook 是一个记录器钩子,故其可使用 log_config

使用方法

cfg.log_config.hooks = [
dict(type='TextLoggerHook'),
dict(type='MMDetWandbHook',
init_kwargs={'project': 'MMDetection-tutorial'},
interval=10,
log_checkpoint=True,
log_checkpoint_metadata=True,
num_eval_images=100)]
MMDetWandbHook 中有四个重要的参数,可帮助您充分利用MMDetection。下面就让我们来具体看看吧。

记录指标

如要记录训练和评估指标以及系统指标,只需使用 init_kwargs 参数即可。该参数可接收键值对字典,然后将其传递给 wandb.init。您可使用它来设置W&B项目名称、设置团队名称(entity,如果您想将运行记录至团队账户的话)、传递配置等等。请在此处查看可传递给 wandb.init 的内容。

训练指标:




验证指标:




系统指标 + 杂项:




设定检查点

MMDetection使用MMCV的 CheckpointHook 来周期性地保存模型检查点。时间周期使用 checkpoint_config.interval 进行定义。不过,这些检查点将保存在本地,所以可能会被新的实验所覆盖。
您可通过使用 log_checkpoint=True 参数来将这些检查点可靠地存储为W&B artifact。在将其保存为 W&B artifact后,您就可轻松地跨机器传输检查点、分别保留不同的模型想法并对变体进行比较。
图 1:W&B artifact用户界面。左侧窗格中显示了不同版本的检查点。您可从“文件”标签页中下载模型,或使用API以编程的方式进行下载。
下面是一些值得注意的地方:
  • 如上所示,用户界面中有3个版本的检查点。这是因为在 checkpoint_config.interval=4 的作用下模型训练了12个epoch。
  • 每个变体都有一个别名 epoch_x,其中的 x 就是epoch编号。
  • 最后的检查点以别名 latest 进行标记。
我们建议您谨慎设置检查点间隔以节省本地和W&B存储空间。

设定带有元数据的检查点

如果检查点中能够带有一些有用的、相关的元数据,那么其实用意义将更加凸显。如果将 log_checkpoint_metadata 设为 True,那么各检查点版本都将包含与其关联的元数据。
元数据的形式是键值对字典,其中包含使用该检查点和当前epoch在验证数据上进行计算的评估度量。
图 2:W&B artifact用户界面。记录的元数据将在“元数据”标签页中显示。
注意,此功能依赖于 CheckpointHookEvalHookDistEvalHook。只有当检查点间隔可被评估间隔整除时才会记录元数据(即,ckpt_interval%eval_interval==0)。

数据集与模型预测可视化

对数据集,尤其是模型预测,进行交互式可视化的能力能够帮助建立和调试出更好的模型。MMDetection使用 EvalHook(单 GPU)与 DistEvalHook(多 GPU)来定期评估验证数据上的模型。评估周期受 evaluation.interval 控制(前文已有提及)。
MMDetWandbHook 的帮助下,您可将验证数据记录为W&B表格,并为模型预测创建多个版本的W&B表格。W&B表格是一种二维的数据网格,每一列都有单一类型的数据——图像、文本、标量等。
MMDetWandbHooknum_eval_images 参数控制着被记录为W&B表格的验证样本数量。
下面是一些值得注意的地方:
  • 如果 num_eval_images=0,那么验证数据与模型预测将不会被记录。
  • 如果 mmdet.core.train_detector API 的 validate=False,那么验证数据与模型预测将不会被记录。
  • 如果 num_eval_images 值大于总的验证样本数量,那么会记录完整的验证数据集。
我们已在上方看到过记录为W&B表格的验证数据。在每个评估间隔,预测边界框和基本事实将被记录为W&B表格。可查看下表中给出的模型预测。



存储方面的简要注意事项

每个评估区间的预测表格都会被记录为唯一版本。注意 media/ 仅有几KB的存储占用。这几KB的产生原因是预测分割mask的预测记录。因此,验证数据只会记录一次,因为W&B表格和后续表会使用对该数据的引用。
图 3:W&B artifact用户界面。在“文件”标签页中您可看到已记录表格的存储占用情况。

使用Weights & Biases进行训练


在本节中,我们将使用MMDetection提供的模块化API来构建训练计划安排。和任何训练计划一样,这里我们需要数据加载器、模型和训练循环。但在开始训练模型之前,我们来看看如何对配置进行记录。

日志配置

如果您使用的是 tools/train.py 来训练模型,那么配置将被自动记录。不过,如果您使用的是自己的训练计划,则请确保按照以下步骤进行操作:
  • 创建一个字典,以 metaexp_name 作为 key,值为配置文件名称。
meta = dict()
meta['exp_name'] = os.path.basename(config_file)
  • • 将配置转储至 work_dir/meta['exp_name'] 目录中。MMDetWandbHook 将自动读取保存的配置文件并对其进行记录。

训练

正如下方的代码片段所示,MMDetection 提供开箱即用的 API,可帮助您构建训练计划的必要组件。
def train(cfg, meta):
# Build dataset
datasets = [mmdet.datasets.build_dataset(cfg.data.train)]

# Build the detector
model = mmdet.models.build_detector(
cfg.model, train_cfg=cfg.get('train_cfg'), test_cfg=cfg.get('test_cfg'))
model.CLASSES = datasets[0].CLASSES

# Dump config for logging
mmcv.mkdir_or_exist(os.path.abspath(cfg.work_dir))
cfg.dump(os.patg.join(cfg.work_dir, meta['exp_name']))

# Train the model.
mmdet.apis.train_detector(model, datasets, cfg, distributed=False, validate=True, meta=meta)
distributed=True 时即可使用 train_detector API来在多显卡设置下训练模型。如果您不想在验证数据集上对模型进行评估,则可设置 validate=False。注意,如果 validateFalse,那么 EvalHookDistEvalHook 将不会初始化。如需进行训练,只需:
train(cfg, meta)
由于我们已在上文中看到过记录的指标、数据和模型预测,下面我们就来看看记录的配置吧。



总结

MMDetection是一种可用于目标检测和分割任务的流行框架。该框架中现在加入了专门的 MMDetWandbHook,这样您就能更充分地利用工作流了。如上所示,其使用与更新配置文件的log_config 一样简单。


希望这一功能能够为您的工作流增添价值。如有任何问题,请随时在下方留下评论,也可在这里这里提交问题。我们还实现了与其他热门目标检测框架(如YOLOv5和PaddleDetection)之间的集成。请在下方查看这些检测框架吧:

推荐阅读


Iterate on AI agents and models faster. Try Weights & Biases today.