Detectron2自定義資料集進行訓練

Detectron2對coco資料格式優先支援。是以在開始之前建議把自己資料修改為标準的coco格式，各種類型資料轉coco格式腳本見：轉換工具箱。

注：這個大佬的資料轉換工具，在box标注那裡貌似會偏移一個像素，不知道修複沒有。但是如果對檢測box的定位精度不是要求很苛刻的話，這個并不會有太大影響。

Detectron2訓練自己資料

1.第一步當然是安裝Detectron2：

系統軟體要求：

• Linux or macOS
• Python >= 3.6
• PyTorch >= 1.3
• torchvision
• OpenCV
• fvcore
• pycocotools
• GCC >= 4.9
• CUDA9/10
• CUDNN

pip install torch torchvision
pip install cython
pip install opencv-python
pip install -U 'git+https://github.com/facebookresearch/fvcore'
pip install pycocotools
pip install pandas matplotlib
           

然後克隆主倉庫：

git clone https://github.com/facebookresearch/detectron2.git
cd detectron2
pip install -e .
           

注： 最後

pip install -e

後面的點不要省略！！！

如果編譯重新安裝，則需要删除build目錄下的内容，即：

cd detectron2
rm -rf build/*
pip install -e .
           

安裝完畢後，在終端執行：

python -c 'import torch; from torch.utils.cpp_extension import CUDA_HOME; print(torch.cuda.is_available(), CUDA_HOME)'
           

如果列印

True

則證明安裝成功！

2.裝完之後訓練自定義的資料集

這一步其實很簡單，任意建立一個py檔案，比如

train.py

:

#-*- coding:utf-8 -*-
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import os
import cv2
import random
import matplotlib.pyplot as plt
import numpy as np

from detectron2.engine import DefaultTrainer
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.data.datasets import register_coco_instances

detectron2_repo_path = "path_to_detectron2/"
img_path = "your_image_path/"
json_path = "your_label_json"


register_coco_instances("mydata", {}, json_path, img_path)
mydata_metadata = MetadataCatalog.get("mydata")
dataset_dicts = DatasetCatalog.get("mydata")


cfg = get_cfg()
cfg.merge_from_file(
    os.path.join(detectron2_repo_path, "configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
)
cfg.DATASETS.TRAIN = ("mydata",)
cfg.DATASETS.TEST = ()  # no metrics implemented for this dataset
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = "detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl"  # initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.02
cfg.SOLVER.MAX_ITER = (
    300
)  # 300 iterations seems good enough, but you can certainly train longer

cfg.MODEL.ROI_HEADS.NUM_CLASSES = 10  # 總共10個類别，不含背景

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
           

上面代碼的

detectron2_repo_path、 img_path和 json_path

分别是你的

Detectron2

倉庫路徑，圖像路徑和json标注檔案路徑。

register_coco_instances("mydata", {}, json_path, img_path)
mydata_metadata = MetadataCatalog.get("mydata")
dataset_dicts = DatasetCatalog.get("mydata")
           

這三行是把你的自定義資料注冊完了，資料類名

mydata

。

cfg.merge_from_file

是加載配置檔案，這裡選用了

mask_rcnn_R_50_FPN_3x.yaml

，可以自行修改。

cfg.MODEL.WEIGHTS

裡直接用MODEL_ZOO的權重檔案初始化了模型。

其他學習率相關的參數自行設定即可，執行

python train.py

即可開始訓練。

可能遇到的問題：

RuntimeError: Not compiled with GPU support:ROIAlign_forward at detectron2/layers/csrc/ROIAlign/ROIAlign.h:73

這個問題解決的方法是保證Torch版本正确，而且cuda/cudnn安裝成功。我解除安裝了pytorch和cuda重裝，然後重新編譯安裝Detectron2，問題就消失了。

參考：

1.how-to-train-detectron2-with-custom-coco-datasets

2.detectron2 GETTING_STARTED