search

facebookresearch / Detectron

FAIR's research platform for object detection research, implementing popular algorithms like Mask R-CNN and RetinaNet.
Apache License 2.0
26.22k stars 5.45k forks source link

Use 1 class dataset, Tain on pre trained model is Error! #939

Closed surserrr closed 5 years ago

surserrr commented 5 years ago

My dataset only has 1 class, and I have changed the NUM_CLASSES:2 in yaml. But when I train on the pre trained model, It is wrong.

INFO net.py: 96: fc6_b [+ momentum] loaded from weights file into gpu_0/fc6_b: (1024,) INFO net.py: 96: fc7_w [+ momentum] loaded from weights file into gpu_0/fc7_w: (1024, 1024) INFO net.py: 96: fc7_b [+ momentum] loaded from weights file into gpu_0/fc7_b: (1024,) INFO net.py: 96: cls_score_w [+ momentum] loaded from weights file into gpu_0/cls_score_w: (81, 1024) Traceback (most recent call last): File "tools/train_net.py", line 138, in main() File "tools/train_net.py", line 120, in main checkpoints = detectron.utils.train.train_model() File "/media/fangjt/3dc196fc-1311-49bb-ad24-5f77d8735211/Detectron-master/detectron/utils/train.py", line 58, in train_model setup_model_for_training(model, weights_file, output_dir) File "/media/fangjt/3dc196fc-1311-49bb-ad24-5f77d8735211/Detectron-master/detectron/utils/train.py", line 174, in setup_model_for_training nu.initialize_gpu_from_weights_file(model, weights_file, gpu_id=0) File "/media/fangjt/3dc196fc-1311-49bb-ad24-5f77d8735211/Detectron-master/detectron/utils/net.py", line 109, in initialize_gpu_from_weights_file src_blobs[src_name].shape) AssertionError: Workspace blob cls_score_w with shape (2, 1024) does not match weights file shape (81, 1024)

The yaml is:

MODEL: TYPE: generalized_rcnn CONV_BODY: FPN.add_fpn_ResNet101_conv5_body NUM_CLASSES: 2 FASTER_RCNN: True NUM_GPUS: 1 SOLVER: WEIGHT_DECAY: 0.0001 LR_POLICY: steps_with_decay BASE_LR: 0.001 GAMMA: 0.1 MAX_ITER: 10000 STEPS: [0, 7000, 9000] FPN: FPN_ON: True MULTILEVEL_ROIS: True MULTILEVEL_RPN: True FAST_RCNN: ROI_BOX_HEAD: fast_rcnn_heads.add_roi_2mlp_head ROI_XFORM_METHOD: RoIAlign ROI_XFORM_RESOLUTION: 7 ROI_XFORM_SAMPLING_RATIO: 2 TRAIN: WEIGHTS: ../fjt-experiments/fastrpn/model_final.pkl DATASETS: ('bingxiang_train',) SCALES: (800,) MAX_SIZE: 1333 BATCH_SIZE_PER_IM: 512 RPN_PRE_NMS_TOP_N: 2000 # Per FPN level TEST: DATASETS: ('bingxiang_val',) SCALE: 800 MAX_SIZE: 1333 NMS: 0.5 RPN_PRE_NMS_TOP_N: 1000 # Per FPN level RPN_POST_NMS_TOP_N: 1000 OUTPUT_DIR: ../fjt-experiments/fastrpn/result

surserrr commented 5 years ago

I solve the problem. If the categories is different from the coco , we can use ImageNet Pretrained Models as the pretrained model. If you want to use End-to-End Pretrained Models, the categories must be 81.