How to finetune the model?

[1960675737]

How to finetune the model when changing the NUM_CLASSES?

[ys0823]

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]

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