以下链接是个人关于detectron2(目标检测框架),所有见解,如有错误欢迎大家指出,我会第一时间纠正。有兴趣的朋友可以加微信:a944284742相互讨论技术。若是帮助到了你什么,一定要记得点赞!因为这是对我最大的鼓励。
detectron2(目标检测框架)无死角玩转-00:目录
前言
通过前面的博客,我们已经知道怎么去评估以及训练一个模型,但是都是基于官方的数据集进行训练,那么我们如何去训练自己的数据呢?那么这篇博客我们来为大家讲解,首先请大家下载好该数据集,如果连接失效,请留言,本人会及时更换:
https://pan.baidu.com/s/1d19d5ke5k779sE-gNeXp8g
下载好之后,本人摆放的位置如下:
1.detectron2-master
......
2.Dataset
--annotations
test_cap.json
train_cap.json
--images
jpg图片
......
放置好之后,那么我来简单的介绍一下该数据集。该数据集主要是一个瓶盖缺陷检测的数据集,总共有10总缺陷,如果加上背景则为11种,打开其中的train_cap.json即可查看,如下:
其中的categories的内容为:
['背景', '瓶盖破损', '瓶盖变形', '瓶盖坏边', '瓶盖断点', '标贴歪斜', '标贴起皱', '标贴气泡', '喷码正常', '喷码异常']
数据的格式,完全按照coco的数据格式,因为这样我们就能直接使用coco的API,可以节约很多时间。
代码编写
参考于tools/train_net.py本人编写了tools/train_my.py程序,如下(直接拷贝即可):
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.data.datasets.coco import load_coco_json
import os
import cv2
import logging
from collections import OrderedDict
import detectron2.utils.comm as comm
from detectron2.utils.visualizer import Visualizer
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.data.datasets.coco import load_coco_json
from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, launch
from detectron2.evaluation import (
CityscapesEvaluator,
COCOEvaluator,
COCOPanopticEvaluator,
DatasetEvaluators,
LVISEvaluator,
PascalVOCDetectionEvaluator,
SemSegEvaluator,
verify_results,
)
import torch
from detectron2.modeling import GeneralizedRCNNWithTTA
import pycocotools
# fmt: off
#CLASS_NAMES = ['背景', '瓶盖破损', '瓶盖变形', '瓶盖坏边', '瓶盖打旋', '瓶盖断点', '标贴歪斜', '标贴起皱', '标贴气泡', '喷码正常', '喷码异常']
# 数据集路径
DATASET_ROOT = '/data/zwh/04.tianchi/01.chongqing/2.Scheme/3.detectron2/2.Dataset'
ANN_ROOT = os.path.join(DATASET_ROOT, 'annotations')
TRAIN_PATH = os.path.join(DATASET_ROOT, 'images')
VAL_PATH = os.path.join(DATASET_ROOT, 'images')
TRAIN_JSON = os.path.join(ANN_ROOT, 'train_cap.json')
VAL_JSON = os.path.join(ANN_ROOT, 'test_cap.json')
# 数据集的子集
PREDEFINED_SPLITS_DATASET = {
"coco_my_train": (TRAIN_PATH, TRAIN_JSON),
"coco_my_test": (VAL_PATH, VAL_JSON),
}
def register_dataset():
"""
purpose: register all splits of dataset with PREDEFINED_SPLITS_DATASET
"""
for key, (image_root, json_file) in PREDEFINED_SPLITS_DATASET.items():
register_dataset_instances(name=key,
json_file=json_file,
image_root=image_root)
def register_dataset_instances(name, json_file, image_root):
"""
purpose: register dataset 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():
DatasetCatalog.register("coco_my_train", lambda: load_coco_json(TRAIN_JSON, TRAIN_PATH))
MetadataCatalog.get("coco_my_train").set(#thing_classes=CLASS_NAMES, # 可以选择开启,但是不能显示中文,所以本人关闭
evaluator_type='coco', # 指定评估方式
json_file=TRAIN_JSON,
image_root=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(VAL_JSON, VAL_PATH))
MetadataCatalog.get("coco_my_val").set(#thing_classes=CLASS_NAMES, # 可以选择开启,但是不能显示中文,所以本人关闭
evaluator_type='coco', # 指定评估方式
json_file=VAL_JSON,
image_root=VAL_PATH)
# 查看数据集标注
def checkout_dataset_annotation(name="coco_my_val"):
#dataset_dicts = load_coco_json(TRAIN_JSON, TRAIN_PATH, name)
dataset_dicts = load_coco_json(TRAIN_JSON, 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
class Trainer(DefaultTrainer):
@classmethod
def build_evaluator(cls, cfg, dataset_name, output_folder=None):
"""
Create evaluator(s) for a given dataset.
This uses the special metadata "evaluator_type" associated with each builtin dataset.
For your own dataset, 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))
elif evaluator_type == "cityscapes":
assert (
torch.cuda.device_count() >= comm.get_rank()
), "CityscapesEvaluator currently do not work with multiple machines."
return CityscapesEvaluator(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 dataset {} 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
def setup(args):
"""
Create configs and perform basic setups.
"""
cfg = get_cfg() # 拷贝default config副本
args.config_file = "../configs/My/retinanet_R_50_FPN_3x.yaml"
cfg.merge_from_file(args.config_file) # 从config file 覆盖配置
cfg.merge_from_list(args.opts) # 从CLI参数 覆盖配置
# 更改配置参数
cfg.DATASETS.TRAIN = ("coco_my_train",) # 训练数据集名称
cfg.DATASETS.TEST = ("coco_my_val",)
cfg.DATALOADER.NUM_WORKERS = 10 # 单线程
cfg.INPUT.CROP.ENABLED = True
cfg.INPUT.MAX_SIZE_TRAIN = 640 # 训练图片输入的最大尺寸
cfg.INPUT.MAX_SIZE_TEST = 640 # 测试数据输入的最大尺寸
cfg.INPUT.MIN_SIZE_TRAIN = (512, 768) # 训练图片输入的最小尺寸,可以吃定为多尺度训练
cfg.INPUT.MIN_SIZE_TEST = 640
cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING = 'range'
cfg.MODEL.RETINANET.NUM_CLASSES = 11 # 类别数
#cfg.MODEL.WEIGHTS = "/homeDocuments/pretrainedModel/Detectron2/R-50.pkl" # 预训练模型权重
cfg.SOLVER.IMS_PER_BATCH = 48 # batch_size=2; iters_in_one_epoch = dataset_imgs/batch_size
# 根据训练数据总数目以及batch_size,计算出每个epoch需要的迭代次数
ITERS_IN_ONE_EPOCH = int(3072 / cfg.SOLVER.IMS_PER_BATCH)
# 指定最大迭代次数
cfg.SOLVER.MAX_ITER = (ITERS_IN_ONE_EPOCH * 100) - 1 # 12 epochs,
# 初始学习率
cfg.SOLVER.BASE_LR = 0.002
# 优化器动能
cfg.SOLVER.MOMENTUM = 0.9
#权重衰减
cfg.SOLVER.WEIGHT_DECAY = 0.0001
cfg.SOLVER.WEIGHT_DECAY_NORM = 0.0
# 学习率衰减倍数
cfg.SOLVER.GAMMA = 0.1
# 迭代到指定次数,学习率进行衰减
cfg.SOLVER.STEPS = (7000,)
# 在训练之前,会做一个热身运动,学习率慢慢增加初始学习率
cfg.SOLVER.WARMUP_FACTOR = 1.0 / 1000
# 热身迭代次数
cfg.SOLVER.WARMUP_ITERS = 1000
cfg.SOLVER.WARMUP_METHOD = "linear"
# 保存模型文件的命名数据减1
cfg.SOLVER.CHECKPOINT_PERIOD = ITERS_IN_ONE_EPOCH - 1
# 迭代到指定次数,进行一次评估
cfg.TEST.EVAL_PERIOD = ITERS_IN_ONE_EPOCH
#cfg.TEST.EVAL_PERIOD = 100
cfg.freeze()
default_setup(cfg, args)
return cfg
def main(args):
cfg = setup(args)
#print(cfg)
# 注册数据集
plain_register_dataset()
# 检测数据集注释是否正确
checkout_dataset_annotation()
# 如果只是进行评估
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 comm.is_main_process():
verify_results(cfg, res)
return res
trainer = Trainer(cfg)
trainer.resume_or_load(resume=args.resume)
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.配置修改,重写tools/train_net.py中的setup(args)函数,主要是为了配置方便一些而已。
2.注册数据集,上面编写了3种方式,分别为
def register_dataset(),
def register_dataset_instances(),
def plain_register_dataset().
本人实验为最后一个
3.重写class Trainer(DefaultTrainer)中的:
def build_evaluator()
def test_with_TTA()
其实,直接从tools/train_net.py从拷贝过来即可
这样就可以了,checkout_dataset_annotation函数,他主要是检测json文件的注释是否正确,打上框的图像图像输出在tools/output下。可以通过该函数验证json文件的正确性。本人训练如下:
额外附加
有的时候,训练过程中可能被中断,我们可能就需要从断点继续训练,在训练的时候,只需要添加 --resume 指令即可,如:
python tools/train_net.py --resume
该指令默认是从last_checkpoint进行训练,如果想从特定的pth训练,把他当作预训练模型即可。还有就是模型评估,和前面你的操作方式一样,这里就不再重复了。