search

HRNet / HRNet-MaskRCNN-Benchmark

Object detection with multi-level representations generated from deep high-resolution representation learning (HRNetV2h).
MIT License
138 stars 37 forks source link

How to finetune the model? #4

Open 1960675737 opened 5 years ago

1960675737 commented 5 years ago

How to finetune the model when changing the NUM_CLASSES?

ys0823 commented 5 years ago

How to finetune the model when changing the NUM_CLASSES?

Maybe you need do like this to delete the last layer for the .pth model

 _d=torch.load(path)
newdict=_d
def removekey(d, listofkeys):
    r=d
    for key in listofkeys:
        del r[key]
    return r

newdict['model'] = removekey(_d['model'], ['module.roi_heads.box.predictor.cls_score.bias','module.roi_heads.box.predictor.cls_score.weight','module.roi_heads.box.predictor.bbox_pred.bias','module.roi_heads.box.predictor.bbox_pred.weight','module.roi_heads.mask.predictor.mask_fcn_logits.weight','module.roi_heads.mask.predictor.mask_fcn_logits.bias'])
torch.save(newdict, 'mymodel.pth')
wondervictor commented 5 years ago

Thanks @ys0823 You can initialize the weights/bias of the output layer in RoIHeads.