Mask R-CNN、Pytorch如何训练自己的数据集

第一次尝试在CSDN上写文章。本人环境：pytorch1.4,cuda10.1,ubantu18.4,这里也可以根据本人配置，在windows上进行配置环境

文章目录

@[TOC]1.配置环境

@[TOC]2.Mask　R-CNN benchmark、cocoapi、apex

@[TOC]3.制作自己的数据集，修改预训练权重，训练自己的数据集

创建虚拟环境：

conda create -n maskrcnn python=3.7
conda activate maskrcnn

下载pytorch、依赖包

这里要注意，pytorch的版本一定要与你的cuda版本完全一致，不然在安装apex时，会报错。

conda install pytorch==1.4.0 torchvision==0.5.0 cudatoolkit=10.1 -c pytorch
conda install ipython pip
pip install ninja yacs cython matplotlib tqdm opencv-python -i https://pypi.tuna.tsinghua.edu.cn/simple

如果下载太慢，可以添加镜像源。

第一步添加镜像站到Anaconda执行如下命令：

conda config --add channels http://mirror.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
conda config --add channels http://mirror.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
conda config --set show_channel_urls yes

第二部还可以附加第三方的conda源：

conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/ 
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/msys2/ 
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/bioconda/ 
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/menpo/ 
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/

安装cocoapi

如果windows系统，这里需要先下载visual studio 2015以上版本。默认安装即可

git clone https://github.com/cocodataset/cocoapi.git
cd cocoapi/PythonAPI
python setup.py build_ext --inplace
python setup.py build_ext install

安装apex

apex一定要是Ubantu系统一定要是git clone,windows可以下载到本地。

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
python setup.py install

如果出现错误fatal error: ATen/cuda/DeviceUtils.cuh: No such file or directory

请参考​​​​​​fatal error: ATen/cuda/DeviceUtils.cuh: No such file or directory_少玩游戏多看代码的博客-CSDN博客

安装maskrcnn-benchmark

git clone https://github.com/facebookresearch/maskrcnn-benchmark.git
cd maskrcnn-benchmark

使用labelme制作coco数据集，这里具体步骤不再赘述

【labelme】13分钟教会你使用labelme的超详细教程_哔哩哔哩_bilibili

关键点：安装好labelme之后，直接在命令行输入labelme，不需要添加标签文本文件

python labelme2cocoAll.py 图片目录 --output 生成的json文件

本人数据集格式

annotation目录下

train2017和val2017一个是训练集，一个是测试集，里面内容是一样的。

 修改预训练权重

本人选择的是e2e_mask_rcnn_R_50_FPN_1x.yaml，复制并更名为my_e2e_mask_rcnn_R_50_FPN_1x.yaml。

修改内容：WEIGHT: "./weights/my_pretrained_R_50.pth"

DATASETS:
TRAIN: ("coco_2017_train", )
TEST: ("coco_2017_val",)

在文本后面添加： OUTPUT_DIR: "./weights/"

将my_e2e_mask_rcnn_R_50_FPN_1x.yaml，复制并更名为

my_test_e2e_mask_rcnn_R_50_FPN_1x.yaml

修改内容：

WEIGHT: "./weights/model_final.pth"

并在maskrcnn benchmark根目录下面新建文件夹weights，训练文件，以及预训练权重，可以私我。

准备预训练文件

python weights/trim_detectron_model.py  --pretrained_path weights/e2e_mask_rcnn_R_50_FPN_1x.pth --save_path weights/my_pretrained_R_50.pth

 修改配置文件之后，需要重新编译。

python setup.py build develop

训练自己的数据集

python tools/train_net.py --config-file configs/my_e2e_mask_rcnn_R_50_FPN_1x.yaml MODEL.ROI_BOX_HEAD.NUM_CLASSES 6 SOLVER.IMS_PER_BATCH 4 SOLVER.BASE_LR 0.001 SOLVER.MAX_ITER 36000 SOLVER.STEPS "(24000, 32000)" TEST.IMS_PER_BATCH 1 MODEL.RPN.FPN_POST_NMS_TOP_N_TRAIN 2000