SOLOV2训练教学

Github: https://github.com/WXinlong/SOLO
我的操作系统是Ubuntu18.04，本文将会分成以下部分：

• 创建数据集
• 修改config
• 模型训练
• 可视化Mask结果
• 模型评估
• 推理预测
• 推理优化

首先安装一些基本库，可参考官方安装说明

$ git clone https://github.com/WXinlong/SOLO.git
$ cd SOLO
$ pip install -r requirements/build.txt
$ pip install “git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI"
$ pip install -v -e .

1.创建数据集

• 准备训练的数据集，使用 labelme 标注数据集并转换为 coco 格式，标注操作可参考：语义分割与实例分割 Labelme标注教学
• 建立训练集和验证集，如下所示
  
• 在 mmdet/datasets 里创建自定义数据。建立my_dataset.py，写入训练的类别。如果标签只有一个，要另外加入 null类别。

from .coco import CocoDataset
from .registry import DATASETS
@DATASETS.register_module
class MyDataset(CocoDataset):
    CLASSES = ['null', 'raccoon']

• 在 mmdet/datasets/__init__.py 加入刚刚定义好的数据集
  

2.修改config

选择要训练的模型架构，所有模型位于 configs 文件夹中。我使用 SOLOv2 最轻量的模型，接着打开 configs/solov2/solov2_light_448_r18_fpn_8gpu_3x.py 修改以下内容:

• num_classes 设定为类别数+1
  
• 更改 dataset_type 和 data_root 路径
  
• 更改 data = dict() 里 train、val、test 的 ann_file、img_prefix 路径
  
• 可修改 optimizer、learning rate、total_epochs 等超参数，其中 work_dir 为存放训练模型的路径。
  

3.模型训练

接下来就开始模型训练啦。。。

python tools/train.py configs/solov2/solov2_light_448_r18_fpn_8gpu_3x.py

若是使用多 GPU （假如是8个GPU）可使用以下指令:

./tools/dist_train.sh configs/solov2/solov2_light_448_r18_fpn_8gpu_3x.py  8

训练好以后会将训练好的模型放在work_dir 设定的文件夹路径下。

4.可视化Mask结果

• 修改tools/test_ins_vis.py 中的 class_names为自定义类别
  

python tools/test_ins_vis.py configs/solov2/solov2_light_448_r18_fpn_8gpu_3x.py weights/solov2_light_release_r18_fpn_8gpu_3x/latest.pth --show --save_dir  work_dirs/vis_solo

• 可视化Mask结果会存放在work_dirs/vis_solo文件夹中
  

5.模型评估

• 评估模型AP
  --eval 参数可以选择要显示哪一个AP
  

python tools/test_ins.py configs/solov2/solov2_light_448_r18_fpn_8gpu_3x.py weights/solov2_light_release_r18_fpn_8gpu_3x/latest.pth --show --out results_solo.pkl --eval segm

若想显示各个类别的AP，可以将mmdet/core/evaluation/coco_utils.py中classwise改为True后，再执行一次命令

python tools/test_ins.py configs/solov2/solov2_light_448_r18_fpn_8gpu_3x.py weights/solov2_light_release_r18_fpn_8gpu_3x/latest.pth --show --out results_solo.pkl --eval segm

6.推理预测

训练好后就可以预测啦，将demo/inference_demo.py 的 config_file 及 checkpoint_file 更改为刚训练好的模型路径：

• 执行预测

python demo/inference_demo.py

运行demo会出现以下警告，但不影响输出结果

• 预测结果
  

7.推理优化

若想对推理进行优化，可以使用ncnn以及TensorRT。
NCNN

• 详细记录solov2的ncnn实现和优化
• https://github.com/DayBreak-u/SOLOV2_ncnn
• solov2模型如何转成NCNN
  TensorRT
• SOLOv2.tensorRT
• TensorRT实现solov2加速
• SOLOV2-TensorRT加速一：可能是第一个实时的高精度实例分割模型

完整的代码地址

参考目录

https://medium.com/ching-i/solov2-%E8%A8%93%E7%B7%B4%E6%95%99%E5%AD%B8-90591960b5c7