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