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