facebook-pytorch maskrcnn-benchmark训练自己的数据集
- 数据准备
gray标签图转coco数据集 https://www.cnblogs.com/yanghailin/p/11215003.html
voc定位数据xml转coco数据集格式json https://www.cnblogs.com/yanghailin/p/11189871.html
我把数据准备成coco格式的。在maskrcnn-benchmark根目录下面的datasets文件夹下,目录结构如下:
└── coco
├── annotations
│ ├── instances_train2014.json
│ └── instances_val2014.json
├── test2014
│ └── cjh_993.png
├── train2014
│ ├── cjh_968.png
│ ├── cjh_969.png
│ ├── cjh_976.png
│ ├── cjh_977.png
│ └── cjh_984.png
└── val2014
├── cjh_985.png
└── cjh_992.png
- 配置文件修改
根目录下的configs文件夹里面有很多yaml文件,我选择的是e2e_mask_rcnn_R_50_FPN_1x.yaml,根据我的需求我更改如下:
MODEL:
META_ARCHITECTURE: "GeneralizedRCNN"
WEIGHT: "catalog://ImageNetPretrained/MSRA/R-50"
BACKBONE:
CONV_BODY: "R-50-FPN"
RESNETS:
BACKBONE_OUT_CHANNELS: 256
RPN:
USE_FPN: True
ANCHOR_STRIDE: (4, 8, 16, 32, 64)
PRE_NMS_TOP_N_TRAIN: 2000
PRE_NMS_TOP_N_TEST: 1000
POST_NMS_TOP_N_TEST: 1000
FPN_POST_NMS_TOP_N_TEST: 1000
ROI_HEADS:
USE_FPN: True
ROI_BOX_HEAD:
POOLER_RESOLUTION: 7
POOLER_SCALES: (0.25, 0.125, 0.0625, 0.03125)
POOLER_SAMPLING_RATIO: 2
FEATURE_EXTRACTOR: "FPN2MLPFeatureExtractor"
PREDICTOR: "FPNPredictor"
#修改成自己任务所需要检测的类别数+1 我是33类 33+1
NUM_CLASSES: 34
ROI_MASK_HEAD:
POOLER_SCALES: (0.25, 0.125, 0.0625, 0.03125)
FEATURE_EXTRACTOR: "MaskRCNNFPNFeatureExtractor"
PREDICTOR: "MaskRCNNC4Predictor"
POOLER_RESOLUTION: 14
POOLER_SAMPLING_RATIO: 2
RESOLUTION: 28
SHARE_BOX_FEATURE_EXTRACTOR: False
MASK_ON: True
DATASETS:
###我用的coco2014数据结构目录####################################
##此处的名字与maskrcnn-benchmark/maskrcnn_benchmark/config/paths_catalog.py 下的名字对应
##"coco_2014_train": {
## "img_dir": "coco/train2014",
## "ann_file": "coco/annotations/instances_train2014.json"
## 和coco数据结构一致的话就不需要修改
TRAIN: ("coco_2014_train","coco_2014_val")
TEST: ("coco_2014_test",)
DATALOADER:
SIZE_DIVISIBILITY: 32
SOLVER:
BASE_LR: 0.02
WEIGHT_DECAY: 0.0001
STEPS: (60000, 80000)
MAX_ITER: 90000
###下面是我新增的##CHECKPOINT_PERIOD:500 训练每迭代500次保存一次模型 #####
## OUTPUT_DIR: 保存模型路径
IMS_PER_BATCH: 3
CHECKPOINT_PERIOD: 500
OUTPUT_DIR: "./my_weights/chejiahao/"
- maskrcnn-benchmark/maskrcnn_benchmark/config/defaults.py 改一处类别数改成自己的
.#大概在214行 就改这一处就可以了 之前是81
_C.MODEL.ROI_BOX_HEAD.NUM_CLASSES = 34
4.训练
在根目录下运行:
python tools/train_net.py --config-file ./demo/my/e2e_mask_rcnn_R_50_FPN_1x.yaml SOLVER.IMS_PER_BATCH 2 SOLVER.BASE_LR 0.0025 SOLVER.MAX_ITER 720000 SOLVER.STEPS "(480000, 640000)" TEST.IMS_PER_BATCH 1 MODEL.RPN.FPN_POST_NMS_TOP_N_TRAIN 2000
就可以训练起来。训练显示的日志:
2019-07-19 17:11:18,959 maskrcnn_benchmark.trainer INFO: Start training
2019-07-19 17:11:30,359 maskrcnn_benchmark.trainer INFO: eta: 4 days, 17:59:34 iter: 20 loss: 2.7949 (3.2260) loss_classifier: 0.9784 (1.3514) loss_box_reg: 0.2458 (0.1996) loss_mask: 0.7680 (0.9672) loss_objectness: 0.6078 (0.6162) loss_rpn_box_reg: 0.0833 (0.0916) time: 0.5245 (0.5700) data: 0.0089 (0.0395) lr: 0.000897 max mem: 2562
2019-07-19 17:11:41,858 maskrcnn_benchmark.trainer INFO: eta: 4 days, 18:29:10 iter: 40 loss: 2.2825 (2.7907) loss_classifier: 0.9786 (1.1715) loss_box_reg: 0.3196 (0.2657) loss_mask: 0.6615 (0.8187) loss_objectness: 0.2881 (0.4529) loss_rpn_box_reg: 0.0693 (0.0820) time: 0.5792 (0.5725) data: 0.0092 (0.0243) lr: 0.000963 max mem: 2656
- 预测
5.1. 修改之前配置的yaml文件e2e_mask_rcnn_R_50_FPN_1x.yaml的WEIGHT: 路径为自己训练好的模型路径/maskrcnn-benchmark/my_weights/chejiahao/model_0003500.pth
MODEL:
META_ARCHITECTURE: "GeneralizedRCNN"
WEIGHT: "maskrcnn-benchmark/my_weights/chejiahao/model_0003500.pth"
BACKBONE:
CONV_BODY: "R-50-FPN"
5.2 修改demo文件夹下面的predictor.py里面的类别标签列表为自己的
# CATEGORIES = [
# "__background",
# "ZhiZaoNianYue",
# "FDJ_PaiL",
# "FDJ_Hao",
# "CJHao",
# ]
在demo文件夹下面新建my_predictor.py
from maskrcnn_benchmark.config import cfg
from predictor import COCODemo
import cv2
import os
config_file = "maskrcnn-benchmark/demo/cjh/e2e_mask_rcnn_R_50_FPN_1x.yaml"
# update the config options with the config file
cfg.merge_from_file(config_file)
# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cuda"])
coco_demo = COCODemo(
cfg,
min_image_size=800,
confidence_threshold=0.7,
)
file_root = 'maskrcnn-benchmark/datasets/coco/test2014/'
file_list = os.listdir(file_root)
save_out = "maskrcnn-benchmark/demo/chejiahao/out/"
for img_name in file_list:
img_path = file_root + img_name
image = cv2.imread(img_path)
predictions = coco_demo.run_on_opencv_image(image)
save_path = save_out + img_name
cv2.imwrite(save_path,predictions)
cv2.namedWindow("img",cv2.WINDOW_NORMAL)
cv2.imshow("img",predictions)
cv2.waitKey(1)
在demo文件夹下运行:
python my_predictor.py
enjoy!
好记性不如烂键盘---点滴、积累、进步!