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