detectron2 DiffusionDet 训练自己的数据集

本文主要是介绍detectron2 DiffusionDet 训练自己的数据集，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

配环境

git clone https://github.com/ShoufaChen/DiffusionDet# 创建环境
conda create -n diffusion python=3.9
conda activate diffusion
conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cudatoolkit=11.3 -c pytorch
pip install opencv-python# 安装detectron2
cd /data2/zy/DiffusionDet/
git clone https://github.com/facebookresearch/detectron2.git
python -m pip install -e detectron2pip install timm # 不装就会报错 No module named 'timm' (diffusion) 

prepare datasets

mkdir -p datasets/coco
mkdir -p datasets/lvisln -s /path_to_coco_dataset/annotations datasets/coco/annotations
ln -s /path_to_coco_dataset/train2017 datasets/coco/train2017
ln -s /path_to_coco_dataset/val2017 datasets/coco/val2017

修改配置文件等

复制一份train_net.py，命名为train.py，在其中添加下列代码注册数据集

#引入以下注释
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.data.datasets.coco import load_coco_json
import pycocotools
#声明类别，尽量保持
CLASS_NAMES =["__background__","Inlet","Slightshort","Generalshort","Severeshort","Outlet"]
# 数据集路径
DATASET_ROOT = '/data2/zy/DiffusionDet/datasets/coco/'
ANN_ROOT = os.path.join(DATASET_ROOT, 'annotations')TRAIN_PATH = os.path.join(DATASET_ROOT, 'train2017')
VAL_PATH = os.path.join(DATASET_ROOT, 'val2017')
TEST_PATH = os.path.join(DATASET_ROOT, 'test2017')TRAIN_JSON = os.path.join(ANN_ROOT, 'instances_train2017.json')
VAL_JSON = os.path.join(ANN_ROOT, 'instances_val2017.json')
TEST_JSON = os.path.join(ANN_ROOT, 'instances_test2017.json')# 声明数据集的子集
PREDEFINED_SPLITS_DATASET = {"coco_my_train": (TRAIN_PATH, TRAIN_JSON),"coco_my_val": (VAL_PATH, VAL_JSON),
}
#===========以下有两种注册数据集的方法，本人直接用的第二个plain_register_dataset的方式 也可以用register_dataset的形式==================
#注册数据集（这一步就是将自定义数据集注册进Detectron2）
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

main中调用注册函数

def main(args):cfg = setup(args)register_dataset() # here to registerif args.eval_only:model = Trainer.build_model(cfg)kwargs = may_get_ema_checkpointer(cfg, model)if cfg.MODEL_EMA.ENABLED:EMADetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR, **kwargs).resume_or_load(cfg.MODEL.WEIGHTS,resume=args.resume)else:DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR, **kwargs).resume_or_load(cfg.MODEL.WEIGHTS,resume=args.resume)res = Trainer.ema_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 restrainer = Trainer(cfg)trainer.resume_or_load(resume=args.resume)return trainer.train()

在 DiffisionDet/configs 下新建demo.yaml，主要是修改batchsize和max_iter

_BASE_: "Base-DiffusionDet.yaml"
MODEL:WEIGHTS: "detectron2://ImageNetPretrained/torchvision/R-50.pkl"RESNETS:DEPTH: 50STRIDE_IN_1X1: FalseDiffusionDet:NUM_PROPOSALS: 100NUM_CLASSES: 5
DATASETS:TRAIN: ("coco_my_train",)TEST:  ("coco_my_val",)
SOLVER:IMS_PER_BATCH: 16BASE_LR: 0.000025STEPS: (5850, 7000)MAX_ITER: 7500# TOTAL_NUM_IMAGES / (IMS_PER_BATCH * NUM_GPUS) * num_epochs = MAX_ITER# 2000/(16*1)*60=7500 
INPUT:MIN_SIZE_TRAIN: (800,)CROP:ENABLED: FalseFORMAT: "RGB"
OUTPUT_DIR: ./OUTPUT/bs16

训练

 python train.py --num-gpus 1     --config-file configs/diffdet.coco.res50.yaml

这篇关于detectron2 DiffusionDet 训练自己的数据集的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！

---