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，問題就消失了。