pytorch 加載部分權重

pytorch加載部分權重

state_list=['rpn.anchor_generator.cell_anchors.0', 'rpn.anchor_generator.cell_anchors.1',
            'rpn.anchor_generator.cell_anchors.2', 'rpn.anchor_generator.cell_anchors.3',
            'rpn.anchor_generator.cell_anchors.4', 'rpn.head.conv.weight', 'rpn.head.conv.bias',
            'rpn.head.cls_logits.weight', 'rpn.head.cls_logits.bias', 'rpn.head.bbox_pred.weight',
            'rpn.head.bbox_pred.bias', 'roi_heads.box.feature_extractor.fc6.weight',
            'roi_heads.box.feature_extractor.fc6.bias', 'roi_heads.box.feature_extractor.fc7.weight',
            'roi_heads.box.feature_extractor.fc7.bias', 'roi_heads.box.predictor.cls_score.weight',
            'roi_heads.box.predictor.cls_score.bias', 'roi_heads.box.predictor.bbox_pred.weight',
            'roi_heads.box.predictor.bbox_pred.bias', 'roi_heads.mask.feature_extractor.mask_fcn1.weight',
            'roi_heads.mask.feature_extractor.mask_fcn1.bias', 'roi_heads.mask.feature_extractor.mask_fcn2.weight',
            'roi_heads.mask.feature_extractor.mask_fcn2.bias', 'roi_heads.mask.feature_extractor.mask_fcn3.weight',
            'roi_heads.mask.feature_extractor.mask_fcn3.bias', 'roi_heads.mask.feature_extractor.mask_fcn4.weight',
            'roi_heads.mask.feature_extractor.mask_fcn4.bias', 'roi_heads.mask.predictor.conv5_mask.weight',
            'roi_heads.mask.predictor.conv5_mask.bias', 'roi_heads.mask.predictor.mask_fcn_logits.weight',
            'roi_heads.mask.predictor.mask_fcn_logits.bias']
state=torch.load(r'D:\PycharmProjects\maskrcnn-benchmark-master\train/e2e_mask_rcnn_X_101_32x8d_FPN_1x_14.pkl')
load_state = {k: v for k, v in state.items() if k not in state_list}

model_state = model.state_dict()

model_state.update(load_state)

model.load_state_dict(state)