detectron2训练自己的数据集（只包含一类数据）

detectron2训练自己的coco格式数据集（只包含一类数据）

detectron2的安装

Linux 或 macOS 或者windows与 Python ≥ 3.6
PyTorch ≥ 1.8 和与PyTorch安装相匹配的 torchvision。在pytorch.org 上将它们一起安装以确保这一点
OpenCV 是可选的，但演示和可视化需要
下面以win系统介绍安装

#虚拟环境的创建
conda create -n detectron2 python=3.7 -y 
conda activate detectron2
#这里选择自己系统对应版本的pytorch 
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
pip install ninja yacs cython matplotlib tqdm
pip install opencv-python==4.4.0.40
# Boundary dependency
pip install scikit-image
pip install shapely
#创建一个项目文件夹例如detectron
mkdir detectron
cd detectron
# install pycocotools. Please make sure you have installed cython.
pip install git+https://github.com/philferriere/cocoapi.git#subdirectory=PythonAPI
# install detectron2
git clone https://github.com/facebookresearch/detectron2.git
python -m pip install -e detectron2

如果pycocotools安装不成功可以去这个地址下载到本地，然后pip install
官方detectron2的文档上说明gcc & g++ ≥ 5.4 是必需的
所以最好安装一下gcc
可以去这个网址下载 下载地址，

一键安装即可
最后打开命令行输入gcc -v 如果有版本号则安装成功

detectron2的训练、评估自己的数据集

你如果需要利用detectron2用到自己的数据集，可以在此detectron2文件内外层先建一个文件夹，例如detectron2/mydatasets，然后放置格式可如下：

接下来有两种方法训练和评估自己的数据集

方法一：创建一个train_net.py文件跑通

在detectron2目录下创建一个train_net.py文件

首先将detectron2/tools/train_net.py的原文件复制粘贴到新建的train_net.py，更改main函数

def main(args):
    cfg = setup(args)
    #coco_my_train是你的数据集自定义的名字，没有要求
    #后面就是你数据集json文件放置和图片放置路径
    #这是训练集
    register_coco_instances("coco_my_train", {}, "mydatasets/coco/annotations/instance_train2017.json",
                            "mydatasets/coco/train2017image")
     #验证集
    register_coco_instances("coco_my_val", {}, "mydatasets/coco/annotations/instances_val2017.json",
                            "mydatasets/coco/val2017image")
"""
这部分代码可以利用注册数据集的元数据得到你的数据字典格式和内容，可以打印和加载注释图片
    custom_metadata = MetadataCatalog.get("coco_my_train")
    dataset_dicts = DatasetCatalog.get("coco_my_train")
    print(dataset_dicts)
    for d in random.sample(dataset_dicts, 1):
        img = cv2.imread(d["file_name"])
        #print(img)
        visualizer = Visualizer(img[:, :, ::-1], metadata=custom_metadata, scale=1)
        vis = visualizer.draw_dataset_dict(d)
        cv2.imshow('Sample',vis.get_image()[:, :, ::-1])
        cv2.waitKey(1500)
"""
    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()

然后更改setup方法
设置训练的模型参数和超参数

def setup(args):
    """
    Create configs and perform basic setups.
    """
    cfg = get_cfg()
    args.config_file = "configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml"
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    # 更改配置参数
    cfg.DATASETS.TRAIN = ("coco_my_train",) # 训练数据集名称
    cfg.DATASETS.TEST = ("coco_my_val",)
    cfg.DATALOADER.NUM_WORKERS = 4  # 单线程

    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，其存在两种配置，分别为 choice 与 range ：
    # range 让图像的短边从 512-768随机选择
    #choice ： 把输入图像转化为指定的，有限的几种图片大小进行训练，即短边只能为 512或者768
    #cfg.INPUT.MIN_SIZE_TRAIN_SAMPLING = 'range'
#  本句一定要看下注释！！！！！！！！
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
    cfg.MODEL.RETINANET.NUM_CLASSES = 1 # 类别数+1（因为有background，也就是你的 cate id 从 1 开始，如果您的数据集Json下标从 0 开始，这个改为您对应的类别就行，不用再加背景类！！！！！）
    #cfg.MODEL.WEIGHTS="/home/yourstorePath/.pth"
    cfg.MODEL.WEIGHTS = "detectron2://ImageNetPretrained/MSRA/R-50.pkl"    # 预训练模型权重
    cfg.SOLVER.IMS_PER_BATCH = 2  # batch_size=2; iters_in_one_epoch = dataset_imgs/batch_size

    # 根据训练数据总数目以及batch_size，计算出每个epoch需要的迭代次数
    #9000为你的训练数据的总数目，可自定义
    ITERS_IN_ONE_EPOCH = int(50/ cfg.SOLVER.IMS_PER_BATCH)

    # 指定最大迭代次数
    cfg.SOLVER.MAX_ITER = 1500，
    # 初始学习率
    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 = (900,)
    # 在训练之前，学习率慢慢增加初始学习率
    cfg.SOLVER.WARMUP_FACTOR = 1.0 / 1000
    # 热身迭代次数
    cfg.SOLVER.WARMUP_ITERS = 500
    #cfg.MODEL.DEVICE = 'gpu'
    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

#训练并评估数据集
python train_net.py

这是评估训练以及评估结果

方法二更改detecron的配置文件训练评估

一. 修改detectron2\data\datasets\builtin.py的文件配置

coco_self_train和coco_self_val就是自己训练数据集的名称,然后是mydataset下的文件路径

并且将builtin.py文件主函数下面的数据集路径换成自己的路径

二. 修改detectron2\data\datasets\builtin_meta.py的文件配置
1.在COCO_CATEGORIES列表下面新建一个自己的类别例如下面我只有一个类

COCO_my_CATEGORIES=[{"color": [220, 20, 60], "isthing": 1, "id": 1, "name": "person"},]

2.新增一个_get_mycoco_instances_meta方法，该方法是用来返还元数据的类实例 类id等

def _get_mycoco_instances_meta():
    thing_ids = [k["id"] for k in COCO_my_CATEGORIES if k["isthing"] == 1]
    thing_colors = [k["color"] for k in COCO_my_CATEGORIES if k["isthing"] == 1]
    assert len(thing_ids) == 1, len(thing_ids)#你有2个类这个数字1就改成2
    # Mapping from the incontiguous COCO category id to an id in [0, 79]
    thing_dataset_id_to_contiguous_id = {k: i for i, k in enumerate(thing_ids)}
    thing_classes = [k["name"] for k in COCO_my_CATEGORIES if k["isthing"] == 1]
    ret = {
        "thing_dataset_id_to_contiguous_id": thing_dataset_id_to_contiguous_id,
        "thing_classes": thing_classes,
        "thing_colors": thing_colors,
    }
    return ret

3.更改 _get_builtin_metadata函数将 return返回改成自己新增的那个方法。
三.使用官方训练命令

cd tools
python train_net.py \--config-file ../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml \
  --num-gpus 1 SOLVER.IMS_PER_BATCH 2 SOLVER.BASE_LR 0.0025

要仅仅评估模型的性能，请使用

python train_net.py \
  --config-file ../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml \
  --eval-only MODEL.WEIGHTS ../output/model_final.pth

添加–resume参数还可以在上次训练终止的地方继续训练，使用上一次的学习率，迭代次数等参数

测试自己的模型的预测效果

cd demo
python demo.py --config-file ../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
  --input input1.jpg  \
  [--other-options]
  --opts MODEL.WEIGHTS ../output/model_fina.pth

其中input.jpg放置于tools目录下即可，也可以创建文件夹例如testimage/input.jpg
[–other-options]参数配置可以参考官方文档。
例如将其可以改成–output result 意思就是测试结果的图片放置在result文件夹下，如果不填任何参数，测试图片会按照yaml原来的参数配置直接显示在opencv窗口

结果

由于计算机资源有限，我只是随便训练了一个10张图片的coco人物数据集，利用生成的模型测试图片，只是做个参考

小伙伴们快去试试吧！