📖 首先需要准备好数据集,这里有labelme标签数据转coco数据集标签的说明:labelme转coco数据集 - 一届书生 - 博客园 (cnblogs.com)

1. 准备工作目录

我们的工作目录,也就是mmdetection目录,如下所示:

|-- configs
|     |--  _base_
|     |---  .......
|-- data
|     |---  coco
|     |     |---  annotations
|     |     |---  train2017
|     |     |---  val2017
|     |     |---  visualization
|-- mmdet
|     |---  core
|     |---  datasets
|     |---  .......
|-- tools
  • configs就是我们的配置文件,里边包含所有的文件。

  • data就是我们的数据集文件,文件目录如上。

  • mmdet是我们所需要修改的目录。

  • tools是我们的mmdetection提供的工具箱,里边包含我们要用的训练和测试文件。

2. 修改mmdetection模型的配置

1️⃣ 如果自己的GPU显存不够用,修改下面文件里的img_scale=(1333, 800),改成小一点的数值。三个文件都要改。

configs/_base_/datasets/coco_detection.py
configs/_base_/datasets/coco_instance.py
configs/_base_/datasets/coco_instance_semantic.py

2️⃣ 选择你要训练的模型对应的配置文件修改,假如我要训练的模型是mask_rcnn_r101_fpn_2x_coco.py,我打开configs/mask_rcnn/mask_rcnn_r101_fpn_2x_coco.py,内容如下:

然后我们找到这个目录下的文件,修改文件里的 num_classes=80 ,修改成自己的类别数目。

configs/_base_/models/mask_rcnn_r50_fpn.py

3️⃣ 修改我们的类别名,两个文件需要修改,第一个文件是:

mmdet/core/evaluation/class_names.py

修改里边的def coco_classes(): ,将return内容修改成自己的类别。

第二个文件:

mmdet/datasets/coco.py

修改里边的class CocoDataset(CustomDataset): ,将 CLASSES = () 修改成自己的类别。

至此,修改结束,我们还需要重新编译一遍,这样才能生效,在我们的mmdetection目录下运行:

python setup.py install

否则会遇到bug:# AssertionError: The `num_classes` (3) in Shared2FCBBoxHead of MMDataParallel does not matches the length of `CLASSES` 80) in CocoDataset - 一届书生 - 博客园 (cnblogs.com)

3. 开始训练

1️⃣ 单GPU训练

python tools/train.py configs/mask_rcnn/mask_rcnn_r101_fpn_2x_coco.py
  • configs/mask_rcnn/mask_rcnn_r101_fpn_2x_coco.py 就是我们要训练的模型

2️⃣ 多GPU训练

bash ./tools/dist_train.sh configs/mask_rcnn/mask_rcnn_r101_fpn_2x_coco.py  2
  • configs/mask_rcnn/mask_rcnn_r101_fpn_2x_coco.py 就是我们要训练的模型
  • 2 是我们的GPU数目

3️⃣ 我们可以刚开始训练便停止,对训练的一些配置进行修改。它会在你的mmdetection目录下自动生成一个work_dirs文件夹,里边包含你模型的配置文件,打开里边的.py文件,例如我的:

mmdetection/work_dirs/mask_rcnn_r101_fpn_2x_coco/mask_rcnn_r101_fpn_2x_coco.py

一般我们进行修改的就是下面这些,官方给的配置文件中所有参数的解释说明:mmdetection-readthedocs-io-zh_CN-latest.pdf

runner = dict(type='EpochBasedRunner', max_epochs=24)
# 最大的epochs,根据自己的情况来调整。
checkpoint_config = dict(interval=1)
# 模型权重的保存的间隔,建议调大一点,否则会保存大量模型权重,占用存储空间,例如interval=8。模型会默认保存最后一次训练的权重
log_config = dict(interval=50, hooks=[dict(type='TextLoggerHook')])
# 日志的输出间隔,建议调小一点,例如interval=4

需要注意的是,修改完配置文件,再训练的时候,训练语句指定的配置文件就是你刚刚修改的了,也就是work_dirs目录下面的。

修改完配置文件后,单 GPU训练

python tools/train.py work_dirs/mask_rcnn_r101_fpn_2x_coco/mask_rcnn_r101_fpn_2x_coco.py

修改完配置文件后,多GPU训练

bash ./tools/dist_train.sh work_dirs/mask_rcnn_r101_fpn_2x_coco/mask_rcnn_r101_fpn_2x_coco.py  2

4. 模型测试

python tools/test.py work_dirs/mask_rcnn_r101_fpn_2x_coco/mask_rcnn_r101_fpn_2x_coco.py work_dirs/mask_rcnn_r101_fpn_2x_coco/latest.pth --show-dir work_dirs/mask_rcnn_r101_fpn_2x_coco/test_show

我是一共有十张图片,7张图片用于训练,3张图片用于测试。有个缺点我没解决,就是螺母的中间,应该为背景,我在用labelme标注过程中都已经标注为_background_,训练的时候,我是用的num_classes=2,我再测试测试。

我自己标注的数据集链接放在这:螺丝螺母数据集.zip - 蓝奏云文件大小:12.0 M|https://wws.lanzouw.com/ixLIX0at5l7i

可视化结果展示:

图1.上图为24epochs,下图为480epochs

同时我们的 work_dirs/mask_rcnn_r101_fpn_2x_coco/ 目录下还会有个json文件,可以可视化我们的一些评价指标的变化情况。

python tools/analysis_tools/analyze_logs.py plot_curve work_dirs/mask_rcnn_r101_fpn_2x_coco/20211015_112915.log.json --keys bbox_mAP segm_mAP

显示结果如下图:

Figure_1

⭐ 文章到此结束了,还有一些小bug,例如上边提到的螺母中间的标注问题,以及最终的评价指标的横坐标epoch显示过于紧密。后续改后对文章继续修改。

Logo

华为开发者空间,是为全球开发者打造的专属开发空间,汇聚了华为优质开发资源及工具,致力于让每一位开发者拥有一台云主机,基于华为根生态开发、创新。

更多推荐