Detectron2: Saving a trained model for inference

[soumik12345]

I used Detectron2 to train a custom object detection model using faster_rcnn_R_50_C4_3x. While I am trying to save the trained model using torch.save or dump in using pickle, I am getting the following error:

AttributeError                            Traceback (most recent call last)
<ipython-input-17-e3c685b650c8> in <module>()
      2 
      3 filehandler = open("hand_detection_faster_rcnn_R_50_C4_3x.pb.pkl", "wb")
----> 4 pickle.dump(predictor, filehandler)
      5 filehandler.close()

AttributeError: Can't pickle local object 'GeneralizedRCNN.__init__.<locals>.<lambda>'

How should I save the detectron2 predictor as a pytorch model for inference?

[Om-Pandey]

Have you tried to save the model using torch.save(net.state_dict(), PATH) and will probably fix your problem.

[soumik12345]

Have you tried to save the model using torch.save(net.state_dict(), PATH) and will probably fix your problem.

The problem is that the network is under abstracted under the Detectron2 Engine, and I am unable to find how to retrieve the network.

[sanatsingh]

Try using build_model function for model using the defined config as stated here :

from detectron2.modeling import build_model
model = build_model(cfg)
torch.save(model.state_dict(), 'checkpoint.pth')

then when loading weights, use :
model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))

--> Map location as required
--> Checkpoint_path the path to checkpoint.pth