Detectron2 部署简单的训练项目
Detectron2 官方文档里的 Getting Started 提供了两种使用 detectron2 的样例。其一是读者大概率已经阅读过的 Colab Notebook ——骑马王子和气球检测，其二是使用命令行执行的 python 文件，包括演示文件 demo.py 及自行用于部署的 train_net.py & plain_train_net.py 。Notebook 已述明使用 Mask-RCNN 进行 mask detection 的简单步骤，包括注册数据集、配置 config 以及训练和验证、可视化结果。由于人力标注费力的原因，自行标注的数据集通常只有目标物体 box 而未标注 mask，因此需使用 Faster-RCNN 进行训练，有关代码可参考 Medium 文章 以及其 代码。
本篇参考 文章 ，简述 python 文件 + 命令行执行的项目部署方法，以便于在 GPU 服务器端运行。
目录一览
本项目部署目录如下：
Project
--configs
----COCO-Detection
faster_rcnn_R_50_FPN_3x.yaml # 官方代码库 copy
Base-RCNN-FPN.yaml # 同上
my_config.yaml # 自己的训练配置文件
--tools
train_net.py # 官方代码库 copy 并自行修改
train.sh
train_resume.sh
eval.sh
--utils
txt2coco.py # 将自己的数据集转化为标准 coco 数据集
数据集目录如下：
MyDataset
--train
01_00001.jpg # 图片名无所谓
...
--val
...
train.json # 位置自行决定
val.json
数据集注册
使用 CocoFormat 的数据集是最优雅的做法，当然你也可以按照官方给定的方法在 tools/rain_net.py 中自定义数据集。关于 csv/voc/labelme 等格式向 coco 格式的转化，可以参考 Github 的代码。本文 utils/txt2coco.py 基于上述代码将自行标注的 txt 格式数据转化为了 coco 格式。
为方便数据集注册，本文借鉴 文章 ，将数据集注册代码包装成类：
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.data.datasets.coco import load_coco_json
class Register:
"""用于注册自己的数据集"""
CLASS_NAMES = ['__background__', 'Red', 'Blue', 'Yellow', 'White', 'Black', 'Other', 'NoHelmet'] # 保留 background 类
ROOT = "/home/your_dataset_dir"
def __init__(self):
self.CLASS_NAMES = Register.CLASS_NAMES or ['__background__', ]
# 数据集路径
self.DATASET_ROOT = Register.ROOT or '/home/yourdir'
# ANN_ROOT = os.path.join(self.DATASET_ROOT, 'COCOformat')
self.ANN_ROOT = self.DATASET_ROOT
self.TRAIN_PATH = os.path.join(self.DATASET_ROOT, 'train')
self.VAL_PATH = os.path.join(self.DATASET_ROOT, 'val')
self.TRAIN_JSON = os.path.join(self.ANN_ROOT, 'train.json')
self.VAL_JSON = os.path.join(self.ANN_ROOT, 'val.json')
# VAL_JSON = os.path.join(self.ANN_ROOT, 'test.json')
# 声明数据集的子集
self.PREDEFINED_SPLITS_DATASET = {
"coco_my_train": (self.TRAIN_PATH, self.TRAIN_JSON),
"coco_my_val": (self.VAL_PATH, self.VAL_JSON),
}
def register_dataset(self):
"""
purpose: register all splits of datasets with PREDEFINED_SPLITS_DATASET
注册数据集（这一步就是将自定义数据集注册进Detectron2）
"""
for key, (image_root, json_file) in self.PREDEFINED_SPLITS_DATASET.items():
self.register_dataset_instances(name=key,
json_file=json_file,
image_root=image_root)
@staticmethod
def register_dataset_instances(name, json_file, image_root):
"""
purpose: register datasets to DatasetCatalog,
register metadata to MetadataCatalog and set attribute
注册数据集实例，加载数据集中的对象实例
"""
DatasetCatalog.register(name, lambda: load_coco_json(json_file, image_root, name))
MetadataCatalog.get(name).set(json_file=json_file,
image_root=image_root,
evaluator_type="coco")
def plain_register_dataset(self):
"""注册数据集和元数据"""
# 训练集
DatasetCatalog.register("coco_my_train", lambda: load_coco_json(self.TRAIN_JSON, self.TRAIN_PATH))
MetadataCatalog.get("coco_my_train").set(thing_classes=self.CLASS_NAMES, # 可以选择开启，但是不能显示中文，这里需要注意，中文的话最好关闭
evaluator_type='coco', # 指定评估方式
json_file=self.TRAIN_JSON,
image_root=self.TRAIN_PATH)
# DatasetCatalog.register("coco_my_val", lambda: load_coco_json(VAL_JSON, VAL_PATH, "coco_2017_val"))
# 验证/测试集
DatasetCatalog.register("coco_my_val", lambda: load_coco_json(self.VAL_JSON, self.VAL_PATH))
MetadataCatalog.get("coco_my_val").set(thing_classes=self.CLASS_NAMES, # 可以选择开启，但是不能显示中文，这里需要注意，中文的话最好关闭
evaluator_type='coco', # 指定评估方式
json_file=self.VAL_JSON,
image_root=self.VAL_PATH)
def checkout_dataset_annotation(self, name="coco_my_val"):
"""
查看数据集标注，可视化检查数据集标注是否正确，
这个也可以自己写脚本判断，其实就是判断标注框是否超越图像边界
可选择使用此方法
"""
# dataset_dicts = load_coco_json(TRAIN_JSON, TRAIN_PATH, name)
dataset_dicts = load_coco_json(self.TRAIN_JSON, self.TRAIN_PATH)
print(len(dataset_dicts))
for i, d in enumerate(dataset_dicts, 0):
# print(d)
img = cv2.imread(d["file_name"])
visualizer = Visualizer(img[:, :, ::-1], metadata=MetadataCatalog.get(name), scale=1.5)
vis = visualizer.draw_dataset_dict(d)
# cv2.imshow('show', vis.get_image()[:, :, ::-1])
cv2.imwrite('out/' + str(i) + '.jpg', vis.get_image()[:, :, ::-1])
# cv2.waitKey(0)
if i == 200:
break
将以上类置于 tools/train_net.py ，并在 main() 函数中第2行添加调用：
def main(args):
cfg = setup(args)
Register().register_dataset() # register my dataset
...
编辑配置文件
从官方下载的 Faster_RCNN 配置文件及其依赖配置最好不要修改，而是基于其配置另写一个文件 configs/my_config.yaml ：
_BASE_: "COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"
DATASETS:
TRAIN: ("coco_my_train",)
TEST: ("coco_my_val",)
MODEL:
RETINANET:
NUM_CLASSES: 8 # 类别数+1, 因为有background
# WEIGHTS: "../tools/output/model_final.pth"
SOLVER:
# IMS_PER_BATCH: 16
# 初始学习率
BASE_LR: 0.00025
# 迭代到指定次数，学习率进行衰减
# STEPS: (210000, 250000)
# MAX_ITER: 270000
CHECKPOINT_PERIOD: 1000
TEST:
EVAL_PERIOD: 3000
如需修改配置，在该文件中修改就好了。
Train & Eval
控制训练和测试的命令行代码置于 shell 文件中更优雅和容易控制。
################ train.sh ################
# Linux 下换行符为 CRLF 的需改为 LF
# lr = 0.00025 * num_gpus
python3 train_net.py \
--config-file ../configs/my_config.yaml \
--num-gpus 4 \
SOLVER.IMS_PER_BATCH 16 \
SOLVER.BASE_LR 0.001 \
SOLVER.MAX_ITER 30000 \
SOLVER.STEPS '(24000, 29000)'
############# train_resume.sh #############
# 断点续 train
# --num-gpus 亲测不能省略
python3 train_net.py \
--config-file ../configs/my_config.yaml \
--num-gpus 4 \
--resume
################# eval.sh #################
python3 train_net.py \
--config-file ../configs/my_config.yaml \
--eval-only \
MODEL.WEIGHTS output/model_final.pth
在命令行中执行 sh train.sh 即可开始训练。
附录
自己的 train_net.py 代码：
#!/usr/bin/env python
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
Detection Training Script.
This scripts reads a given config file and runs the training or evaluation.
It is an entry point that is made to train standard models in detectron2.
In order to let one script support training of many models,
this script contains logic that are specific to these built-in models and therefore
may not be suitable for your own project.
For example, your research project perhaps only needs a single "evaluator".
Therefore, we recommend you to use detectron2 as an library and take
this file as an example of how to use the library.
You may want to write your own script with your datasets and other customizations.
"""
import logging
import os
from collections import OrderedDict
import cv2
import torch
import detectron2.utils.comm as comm
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.data.datasets import load_coco_json
from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, hooks, launch
from detectron2.evaluation import (
CityscapesInstanceEvaluator,
CityscapesSemSegEvaluator,
COCOEvaluator,
COCOPanopticEvaluator,
DatasetEvaluators,
LVISEvaluator,
PascalVOCDetectionEvaluator,
SemSegEvaluator,
verify_results,
)
from detectron2.modeling import GeneralizedRCNNWithTTA
from detectron2.utils.visualizer import Visualizer
class Trainer(DefaultTrainer):
"""
We use the "DefaultTrainer" which contains pre-defined default logic for
standard training workflow. They may not work for you, especially if you
are working on a new research project. In that case you can write your
own training loop. You can use "tools/plain_train_net.py" as an example.
"""
@classmethod
def build_evaluator(cls, cfg, dataset_name, output_folder=None):
"""
Create evaluator(s) for a given datasets.
This uses the special metadata "evaluator_type" associated with each builtin datasets.
For your own datasets, you can simply create an evaluator manually in your
script and do not have to worry about the hacky if-else logic here.
"""
if output_folder is None:
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
evaluator_list = []
evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
if evaluator_type in ["sem_seg", "coco_panoptic_seg"]:
evaluator_list.append(
SemSegEvaluator(
dataset_name,
distributed=True,
num_classes=cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES,
ignore_label=cfg.MODEL.SEM_SEG_HEAD.IGNORE_VALUE,
output_dir=output_folder,
)
)
if evaluator_type in ["coco", "coco_panoptic_seg"]:
evaluator_list.append(COCOEvaluator(dataset_name, cfg, True, output_folder))
if evaluator_type == "coco_panoptic_seg":
evaluator_list.append(COCOPanopticEvaluator(dataset_name, output_folder))
if evaluator_type == "cityscapes_instance":
assert (
torch.cuda.device_count() >= comm.get_rank()
), "CityscapesEvaluator currently do not work with multiple machines."
return CityscapesInstanceEvaluator(dataset_name)
if evaluator_type == "cityscapes_sem_seg":
assert (
torch.cuda.device_count() >= comm.get_rank()
), "CityscapesEvaluator currently do not work with multiple machines."
return CityscapesSemSegEvaluator(dataset_name)
elif evaluator_type == "pascal_voc":
return PascalVOCDetectionEvaluator(dataset_name)
elif evaluator_type == "lvis":
return LVISEvaluator(dataset_name, cfg, True, output_folder)
if len(evaluator_list) == 0:
raise NotImplementedError(
"no Evaluator for the datasets {} with the type {}".format(
dataset_name, evaluator_type
)
)
elif len(evaluator_list) == 1:
return evaluator_list[0]
return DatasetEvaluators(evaluator_list)
@classmethod
def test_with_TTA(cls, cfg, model):
logger = logging.getLogger("detectron2.trainer")
# In the end of training, run an evaluation with TTA
# Only support some R-CNN models.
logger.info("Running inference with test-time augmentation ...")
model = GeneralizedRCNNWithTTA(cfg, model)
evaluators = [
cls.build_evaluator(
cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, "inference_TTA")
)
for name in cfg.DATASETS.TEST
]
res = cls.test(cfg, model, evaluators)
res = OrderedDict({k + "_TTA": v for k, v in res.items()})
return res
class Register:
"""用于注册自己的数据集"""
CLASS_NAMES = ['__background__', 'Red', 'Blue', 'Yellow', 'White', 'Black', 'Other', 'NoHelmet']
ROOT = "/home/mydir"
def __init__(self):
self.CLASS_NAMES = Register.CLASS_NAMES or ['__background__', ]
# 数据集路径
self.DATASET_ROOT = Register.ROOT or '/home/yourdir'
# ANN_ROOT = os.path.join(self.DATASET_ROOT, 'COCOformat')
self.ANN_ROOT = self.DATASET_ROOT
self.TRAIN_PATH = os.path.join(self.DATASET_ROOT, 'train')
self.VAL_PATH = os.path.join(self.DATASET_ROOT, 'val')
self.TRAIN_JSON = os.path.join(self.ANN_ROOT, 'train.json')
self.VAL_JSON = os.path.join(self.ANN_ROOT, 'val.json')
# VAL_JSON = os.path.join(self.ANN_ROOT, 'test.json')
# 声明数据集的子集
self.PREDEFINED_SPLITS_DATASET = {
"coco_my_train": (self.TRAIN_PATH, self.TRAIN_JSON),
"coco_my_val": (self.VAL_PATH, self.VAL_JSON),
}
def register_dataset(self):
"""
purpose: register all splits of datasets with PREDEFINED_SPLITS_DATASET
注册数据集（这一步就是将自定义数据集注册进Detectron2）
"""
for key, (image_root, json_file) in self.PREDEFINED_SPLITS_DATASET.items():
self.register_dataset_instances(name=key,
json_file=json_file,
image_root=image_root)
@staticmethod
def register_dataset_instances(name, json_file, image_root):
"""
purpose: register datasets to DatasetCatalog,
register metadata to MetadataCatalog and set attribute
注册数据集实例，加载数据集中的对象实例
"""
DatasetCatalog.register(name, lambda: load_coco_json(json_file, image_root, name))
MetadataCatalog.get(name).set(json_file=json_file,
image_root=image_root,
evaluator_type="coco")
def plain_register_dataset(self):
"""注册数据集和元数据"""
# 训练集
DatasetCatalog.register("coco_my_train", lambda: load_coco_json(self.TRAIN_JSON, self.TRAIN_PATH))
MetadataCatalog.get("coco_my_train").set(thing_classes=self.CLASS_NAMES, # 可以选择开启，但是不能显示中文，这里需要注意，中文的话最好关闭
evaluator_type='coco', # 指定评估方式
json_file=self.TRAIN_JSON,
image_root=self.TRAIN_PATH)
# DatasetCatalog.register("coco_my_val", lambda: load_coco_json(VAL_JSON, VAL_PATH, "coco_2017_val"))
# 验证/测试集
DatasetCatalog.register("coco_my_val", lambda: load_coco_json(self.VAL_JSON, self.VAL_PATH))
MetadataCatalog.get("coco_my_val").set(thing_classes=self.CLASS_NAMES, # 可以选择开启，但是不能显示中文，这里需要注意，中文的话最好关闭
evaluator_type='coco', # 指定评估方式
json_file=self.VAL_JSON,
image_root=self.VAL_PATH)
def checkout_dataset_annotation(self, name="coco_my_val"):
"""
查看数据集标注，可视化检查数据集标注是否正确，
这个也可以自己写脚本判断，其实就是判断标注框是否超越图像边界
可选择使用此方法
"""
# dataset_dicts = load_coco_json(TRAIN_JSON, TRAIN_PATH, name)
dataset_dicts = load_coco_json(self.TRAIN_JSON, self.TRAIN_PATH)
print(len(dataset_dicts))
for i, d in enumerate(dataset_dicts, 0):
# print(d)
img = cv2.imread(d["file_name"])
visualizer = Visualizer(img[:, :, ::-1], metadata=MetadataCatalog.get(name), scale=1.5)
vis = visualizer.draw_dataset_dict(d)
# cv2.imshow('show', vis.get_image()[:, :, ::-1])
cv2.imwrite('out/' + str(i) + '.jpg', vis.get_image()[:, :, ::-1])
# cv2.waitKey(0)
if i == 200:
break
def setup(args):
"""
Create configs and perform basic setups.
"""
cfg = get_cfg()
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
default_setup(cfg, args)
return cfg
def main(args):
cfg = setup(args)
Register().register_dataset() # register my dataset
if args.eval_only:
model = Trainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume
)
res = Trainer.test(cfg, model)
if cfg.TEST.AUG.ENABLED:
res.update(Trainer.test_with_TTA(cfg, model))
if comm.is_main_process():
verify_results(cfg, res)
return res
"""
If you'd like to do anything fancier than the standard training logic,
consider writing your own training loop (see plain_train_net.py) or
subclassing the trainer.
"""
trainer = Trainer(cfg)
trainer.resume_or_load(resume=args.resume)
if cfg.TEST.AUG.ENABLED:
trainer.register_hooks(
[hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))]
)
return trainer.train()
if __name__ == "__main__":
args = default_argument_parser().parse_args()
print("Command Line Args:", args)
launch(
main,
args.num_gpus,
num_machines=args.num_machines,
machine_rank=args.machine_rank,
dist_url=args.dist_url,
args=(args,),
)
参考
[1]Getting Started with Detectron2 - Github
[2]利用detectron2快速使用faster RCNN - Medium
[3]Detectron2训练自己的数据集（较详细） - CSDN
