SOLOV2训练教学

2023-09-25 12:51:15

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