之前的文章有介紹過 YOLACT、YolactEdge，詳細可參考: YOLACT (You Only Look At CoefficienTs) 系列介紹，本文將要來介紹如何進行訓練。
YOLACT github: https://github.com/dbolya/yolact
YolactEdge github: https://github.com/haotian-liu/yolact_edge
我的作業系統是 Ubuntu 18.04.5，版本為 CUDA 10.2, cuDNN 7.6.5, python 3.7.9, torch 1.6.0, torchvision 0.7.0
本文將會分為三部分: 準備資料集、YOLACT 訓練、YolactEdge 訓練
準備資料集
使用 labelme 做標記並轉換為 coco 格式，標記操作可參考: Segmentation Label 標註教學，轉換後會呈現下列樣子
接著把 JPEGImages, annotations.json 放到 yolact/data 或 yolact_edge/data，yolact 及 yolact_edge 是 git clone 的檔案。
YOLACT 訓練
- 首先安裝需要的依賴及下載整包 yolact
$ pip install cython$ pip install opencv-python pillow pycocotools matplotlib$ git clone https://github.com/dbolya/yolact.git$ cd yolact
- 下載 pretrained weight，我使用 Resnet50 作為 backbone
- 建立 weights 資料夾，把剛剛下載的權重放入資料夾內
- 修改 data/config.py
在 dataset_base = Config() 後加入以下內容。由於我的類別只有一個 (person)，因此 label_map 為 {1:1}
my_custom_dataset = dataset_base.copy({'name': 'My Dataset','train_images': './data/person/','train_info': './data/person/annotations.json','valid_images': './data/person/','valid_info': './data/person/annotations.json','has_gt': True,'class_names': ('person'),'label_map': {1:1}})
修改 yolact_base_config = coco_base_config()，將 dataset 改為剛剛設定的 dataset 名稱、num_classes 改為 background + 類別數、backbone 的部分，我使用 Resnet50，因此改為 resnet50_backbone，其餘的可以根據自己需求做更改。
yolact_base_config = coco_base_config.copy({'name': 'yolact_base',# Dataset stuff
'dataset': my_custom_dataset,'num_classes': 2,# Image Size
'max_size': 550,# Training params
'lr_steps': (2800, 6000, 7000, 7500),'max_iter': 8000,# Backbone Settings
'backbone': resnet50_backbone.copy({
'selected_layers': list(range(1, 4)),
'use_pixel_scales': True,
'preapply_sqrt': False,
'use_square_anchors': True, # This is for backward compatability with a bug'pred_aspect_ratios': [ [[1, 1/2, 2]] ]*5,
'pred_scales': [[24], [48], [96], [192], [384]],
}),
....})
- 訓練模型
接著就可以來訓練模型啦~~ 訓練完的模型會儲存在 weights 資料夾
$ python train.py --config=yolact_my_config
若出現 RuntimeError: Expected a ‘cuda’ device type for generator but found ‘cpu’
→ 更改第 249 行 data_loader = data.DataLoader() 裡 shuffle=False
參考資料: https://github.com/dbolya/yolact/issues/664
- 評估預測
$ python eval.py --trained_model=weights/yolact_base_3999_8000.pth --score_threshold=0.3 --top_k=100 --image=test.jpg
YolactEdge 訓練
YolactEdge 和 Yolact 的訓練方式差不多，差別在於 YolactEdge 可以使用 TensorRT inference。
- 首先要安裝 TensorRT，安裝教學可參考: TensorRT 介紹與安裝教學
- 安裝 torch2trt
$ git clone https://github.com/NVIDIA-AI-IOT/torch2trt$ cd torch2trt$ sudo python setup.py install --plugins
- 安裝需要的依賴及整包 YolactEdge
$ pip install cython$ pip install opencv-python pillow matplotlib$ pip install git+https://github.com/haotian-liu/cocoapi.git#"egg=pycocotools&subdirectory=PythonAPI"$ pip install GitPython termcolor tensorboard$ git clone https://github.com/haotian-liu/yolact_edge.git$ cd yolact_edge
- 下載 coco calib_images → 🔍，放入 data 資料夾內
- 下載 pretrained weight 後，我使用 MobileNetV2 作為 backbone
- 建立 weights 資料夾，把剛剛下載的權重放入資料夾內
- 修改 data/config.py
在 dataset_base = Config() 後加入以下內容
my_custom_dataset = dataset_base.copy({'name': 'My Dataset','train_images': './data/person/','train_info': './data/person/annotations.json','valid_images': './data/person/','valid_info': './data/person/annotations.json','has_gt': True,'class_names': ('person'),'label_map': {1:1}})
修改 yolact_edge_mobilenetv2_config = yolact_edge_config()，參數根據自己的需求做更改。
yolact_edge_mobilenetv2_config = yolact_edge_config.copy({'name': 'yolact_edge_mobilenetv2','torch2trt_max_calibration_images': 100,'dataset': my_custom_dataset,'num_classes': 2,'max_size': 550,'lr_steps': (2800, 6000, 7000, 7500),'max_iter': 8000,'use_tensorrt_safe_mode': False,'backbone': mobilenetv2_backbone})
- 訓練模型
接著就可以來訓練模型啦~~ 訓練完的模型會儲存在 weights 資料夾
$ python train.py --config=yolact_edge_mobilenetv2_config --save_interval=100
- 評估預測
預設是使用 TensorRT INT8，若不使用 TensorRT 或是要使用 TensorRT FP16
# 預設使用 TensorRT INT8
$ python eval.py --trained_model=weights/yolact_edge_mobilenetv2_3999_8000.pth --score_threshold=0.3 --top_k=100 --image=test.jpg# 不使用 TensorRT
$ python eval.py --trained_model=weights/yolact_edge_mobilenetv2_3999_8000.pth --score_threshold=0.3 --top_k=100 --image=test.jpg --disable_tensorrt# 使用 TensorRT FP16
$ python eval.py --trained_model=weights/yolact_edge_mobilenetv2_3999_8000.pth --score_threshold=0.3 --top_k=100 --image=test.jpg --use_fp16_tensorrt
若 TensorRT 在轉換時遇到問題，可使用作者設置的安全模式
$ python eval.py --trained_model=weights/yolact_edge_mobilenetv2_3999_8000.pth --score_threshold=0.3 --top_k=100 --image=test.jpg --use_tensorrt_safe_mode
Segmentation 相關文章
影像分割 Image Segmentation — 語義分割 Semantic Segmentation(1)
