ubuntu16.04使用官方mask-rcnn (caffe2实现)训练自己的数据集

最新推荐文章于 2025-05-26 17:07:00 发布

橘的皮瓜着披。

最新推荐文章于 2025-05-26 17:07:00 发布

阅读量1.9k

点赞数

CC 4.0 BY-SA版权

分类专栏： mask-rcnn

本文链接：https://blog.youkuaiyun.com/lfs666666/article/details/88129352

本文详细介绍了如何在Ubuntu16.04上使用Facebook的Detectron框架，基于Caffe2实现Mask-RCNN训练自己的数据集。内容涵盖源代码和预训练模型获取、环境搭建、数据集准备、注册新数据集、修改配置文件以适应Mask-RCNN训练，并记录了训练过程中遇到的错误及解决方案。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1、源代码及预训练模型准备
(1)官方源码下载地址:
https://github.com/facebookresearch/Detectron
(2)预训练模型下载地址(R-50):
https://dl.fbaipublicfiles.com/detectron/ImageNetPretrained/MSRA/R-50.pkl

2、运行环境搭建

略，详见源码的md说明文件。。。

3、数据集准备
(1)数据集存放位置

$DETECTRON/detectron/datasets/data

把数据集放置在如上目录下.

(2)数据集目录结构(当然这个目录结构可以自由设定),比如我的:

hainu
|_ train
|  |_ <im-1-name>.jpg
|  |_ ...
|  |_ <im-N-name>.jpg
|_ val
|  |_ <im-1-name>.jpg
|  |_ ...
|  |_ <im-N-name>.jpg
|_ annotations
   |_ instances_train.json
   |_ instances_val.json

(3)注册自己的数据集

在$DETECTRON//detectron/datasets/dataset_catalog.py文件中注册自己的数据集

打开dataset_catalog.py可看到coco_2014数据集注册方式如下:

    'coco_2014_train': {
        _IM_DIR:
            _DATA_DIR + '/coco/coco_train2014',
        _ANN_FN:
            _DATA_DIR + '/coco/annotations/instances_train2014.json'
    },
    'coco_2014_val': {
        _IM_DIR:
            _DATA_DIR + '/coco/coco_val2014',
        _ANN_FN:
            _DATA_DIR + '/coco/annotations/instances_val2014.json'
    },

照猫画虎注册自己的数据集:

    'anthrax_train': {
        _IM_DIR:
            _DATA_DIR + '/hainu/train',
        _ANN_FN:
            _DATA_DIR + '/hainu/annotations/instance_train.json'
    },
    'anthrax_val': {
        _IM_DIR:
            _DATA_DIR + '/hainu/val',
        _ANN_FN:
            _DATA_DIR + '/hainu/annotations/instance_val.json'
    }

至此,数据集注册完成.

4、开始用自己的数据集训练mask-rcnn
注：训练及测试网络命令：

#训练
python2 tools/train_net.py \
    --cfg configs/getting_started/tutorial_1gpu_e2e_mask_rcnn_R-50-FPN.yaml \
    OUTPUT_DIR  /home/output
#预测有test_net.py和infer_simple.py两种，其中test_net.py有输出详细的评价指标而infer_simple.py没有。即预测时可以：
#预测
python2 tools/test_net.py \
    --cfg configs/12_2017_baselines/e2e_mask_rcnn_R-50-FPN_2x.yaml \
    TEST.WEIGHTS /home/output/model_final.pkl \
    NUM_GPUS 1

#也可以这样预测，此方式把要测试的图片放到$DETECTRON/demo目录下，下述wts参数请自行修改
python tools/infer_simple.py \
    --cfg configs/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml \
    --output-dir /tmp/detectron-visualizations \
    --image-ext jpg \
    --wts https://dl.fbaipublicfiles.com/detectron/35861858/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml.02_32_51.SgT4y1cO/output/train/coco_2014_train:coco_2014_valminusminival/generalized_rcnn/model_final.pkl \
    demo

(1)训练前修改$DETECTRON/detectron/configs/getting_started目录下的yaml文件

(2)2019年2月份下载下来的源码中,发现在$DETECTRON/detectron/configs/getting_started目录下的yaml文件都只是适用于训练faster-rcnn的,因此要训练mask-rcnn需要把mask部分添加进去（下面有说明）.

###########################出错记录####