Dec 11, 2020

使用 mmdectection 简单训练一个 faster-rcnn 模型

今天整理下怎么简单使用 mmdetection (使用 pytorch) 训练一个 Faster-RCNN 模型，笔者使用的是一个 COCO 格式的数据集。

一、准备环境及仓库

  1. 创建开发环境

    # 创建新的 conda 环境
conda create -n open-mmlab python=3.7 -y
# 激活 conda 环境
conda activate open-mmlab

  2. 安装 pytorch

    conda install pytorch==1.6.0 torchvision==0.7.0 cudatoolkit=10.1 -c pytorch -y

  3. 安装 mmcv-full

    pip install https://download.openmmlab.com/mmcv/dist/latest/torch1.6.0/cu101/mmcv_full-latest%2Btorch1.6.0%2Bcu101-cp37-cp37m-manylinux1_x86_64.whl

  4. clone 仓库

    git clone https://github.com/open-mmlab/mmdetection.git
cd mmdetection

  5. 安装仓库中开发需要的依赖

    pip install -r requirements/build.txt
pip install -v -e .  # or "python setup.py develop"

二、准备数据集

下载 coco 格式的数据集，或者自己标注制作 coco 格式的数据集，比如我们的数据集的目录是 /path/to/your/dataset（绝对路径），需要在 mmdetection 中建立 data 目录，在其下面建立数据集目录的软链接：

# 此时在 mmdetection 目录下，建立 data 目录，如果已经存在可以省略此步
mkdir data
# 建立数据集软链接
ln -s /path/to/your/dataset data/mydetaset

建立软链接时使用的目录路径必须是绝对路径

看一下目录结构，这里是 coco 格式，都差不多，以我手头的一个数据集为例，目录结构如下：

$ tree -L 2 data/mydetaset                          
data/mydetaset
├── annotations
│   ├── test_image_info.json
│   ├── train1.json
│   ├── train2.json
│   ├── val1.json
│   └── val2.json
└── images
    ├── test
    ├── train1
    ├── train2
    ├── val1
    └── val2

7 directories, 5 files

三、修改配置文件

这里以 configs/faster_rcnn/faster_rcnn_x101_64x4d_fpn_2x_coco.py 为例，直接修改配置适配训练任务。

打开这个配置文件发现是引用其它的配置文件，一步步的找到相应的配置文件进行修改：

  1. configs/base/models/faster_rcnn_r50_fpn.py

    这里主要修改检测目标类型个数，比如我修改为 2 个，修改 46 行的 num_classes 的值为 2 即可。

  2. configs/base/schedules/schedule_2x.py

    这里我要使用一张显卡进行训练，所以这里需要调整学习率为原来的 $\frac{1}{8}$ ，修改第 2 行的 lr 即可，修改为 $0.025$。

    原来的值一般是针对 8 卡训练的，因为这里使用一张卡，所以这里调整为原来的 $\frac{1}{8}$

  3. configs/base/datasets/coco_detection.py

    需要修改数据集相关的配置，主要是路径及取样数

    • data_root：这里要修改为上面我们创建的数据集的软链接，为 data/mydetaset/

      这里的目录最后要带着 /

    • samples_per_gpu：如果显卡显存有限，这里就调成 1 吧

    • train：需要对应数据集修改 img_prefix 和 ann_file，以上边准备的数据集为例，修改标注文件和图像目录，同样需要注意最后带上 /:

      train=dict(
  type=dataset_type,
  ann_file=data_root + 'annotations/train1.json',
  img_prefix=data_root + 'images/train1/',
  pipeline=train_pipeline),

    • val：需要对应数据集修改 img_prefix 和 ann_file，以上边准备的数据集为例，修改标注文件和图像目录，同样需要注意最后带上 /:

      val=dict(
  type=dataset_type,
  ann_file=data_root + 'annotations/val1.json',
  img_prefix=data_root + 'images/val1/',
  pipeline=test_pipeline),

    • train：需要对应数据集修改 img_prefix 和 ann_file，以上边准备的数据集为例，修改标注文件和图像目录，同样需要注意最后带上 /:

      val=dict(
  type=dataset_type,
  ann_file=data_root + 'annotations/val1.json',
  img_prefix=data_root + 'images/val1/',
  pipeline=test_pipeline),

四、修改代码中的问题

加载数据集的时候，可以修改数据集加载器类中的代码，实现自动加载正确的类别信息及所有的数据，否则会无法成加载数据倒是训练失败。

修改文件 mmdet/datasets/coco.py 中初始化数据中 categories 的加载实现，很简单，只需要去掉 self.coco.get_cat_ids() 方法中的传参，第 58 行代码修改为：

self.cat_ids = self.coco.get_cat_ids()

另外，保证一致性，还是修改一下 self.CLASSES吧，按照正确顺序写上所有类型即可，这里就不展示了。

五、训练

使用工具脚本 tools/train.py 即可开启一个简单的训练任务，记得指定 gpu：

python tools/train.py configs/faster_rcnn/faster_rcnn_x101_64x4d_fpn_2x_coco.py --gpu-ids 1 --work-dir output
  • configs/faster_rcnn/faster_rcnn_x101_64x4d_fpn_2x_coco.py：指定使用的配置，就是上面决定使用的那个
  • –gpu-ids：指定使用 1 号卡
  • –work-dir：指定模型训练日志及中间 epoch 文件或其它文件的存储目录为 output
访问量