之前的文章有介紹過 SOLO、SOLOv2，詳細可參考: SOLO 與 SOLOv2 論文閱讀，本文將要來介紹如何進行訓練。
🔖 Github: https://github.com/WXinlong/SOLO
我的作業系統是 Ubuntu 18.04.5，版本為 CUDA 10.2, cuDNN 7.6.5, python 3.7.9, torch 1.9.0, torchvision 0.10.0
本文將會分為以下部分:
- 創建數據集
- 修改 config
- 模型訓練
- 視覺化 Mask 結果
- 模型評估
- 推理預測
- Inference Optimization
首先需安裝一些 library，可參考官方安裝說明→ 🔍
$ 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 .
Code 和 dataset 會附在文章最後~
創建數據集
- 準備要訓練的資料集，使用 labelme 做標記並轉換為 coco 格式，標記操作可參考: Segmentation Label 標註教學
# 轉換為 coco 格式
$ cd labelme/examples/instance_segmentation$ python labelme2coco.py <data> <data_output> --labels <label.txt path>
- 建立 data 資料夾放入數據集，資料放置如下
- 在 mmdet/datasets 裡創建自定義數據集。建立 my_dataset.py，寫入訓練的類別。若 label 只有一個，則要另外加入 ‘null’ 類別。
from .coco import CocoDataset
from .registry import DATASETS@DATASETS.register_module
class MyDataset(CocoDataset):
CLASSES = [‘null’, ‘raccoon’]
- 在 mmdet/datasets/__init__.py 加入剛剛定義好的數據集
修改 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 為存放訓練模型的路徑。
模型訓練
接下來就可以開始訓練模型啦~
$ python tools/train.py configs/solov2/solov2_light_448_r18_fpn_8gpu_3x.py若是使用多顆 GPU 可使用以下指令
./tools/dist_train.sh configs/solov2/solov2_light_448_r18_fpn_8gpu_3x.py 8訓練好後，會儲存至 work_dir 設定的資料夾路徑
視覺化 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 資料夾中
模型評估
- 評估模型 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
推理預測
訓練好後就可以進行預測啦~ 將 demo/inference_demo.py 的 config_file 及 checkpoint_file 更改為剛訓練好的模型路徑
- 執行預測
$ python demo/inference_demo.py運行 demo 會出現以下警告，但不影響輸出結果
- 預測結果
Inference Optimization
若要將推理運算進行優化，可以使用 NCNN 及 TensorRT
