How to save model to MODEL ZOO pkl format

[HaowenGuan]

I have trained a R50-FPN with customized dataset and I hope to save the model to .pkl file just like the pretrained weights in MODEL_ZOO.md.
However, after save my model to .pkl with the code in #3535, my pretrained weight does not work the same as pretrained weight in MODEL_ZOO.md.

I inspect the model zoo format, the model is in OrderedDict format, but checkpointer model is different, how can we do the convertion?

What is the trick to save a model as pkl file just like MODEL_ZOO.md?
Is there any function in detectron2 can help this?

[HaowenGuan]

Well, I figured it out, and I am gonna answer myself.

    args = default_argument_parser().parse_args()

    cfg = setup(args)

    model = Trainer.build_model(cfg)
    DetectionCheckpointer(model).resume_or_load(cfg.MODEL.WEIGHTS, resume=False)
    checkpointer = DetectionCheckpointer(model, save_dir='output/supervised')

    from collections import OrderedDict
    weights = OrderedDict()
    for p in checkpointer.model.state_dict():
        weights[p] = checkpointer.model.state_dict()[p].numpy()

    with open(os.path.join("output/model_supervised.pkl"), 'wb') as f:
        myModel = {'model': weights, '__author__': "youself"}
        pickle.dump(myModel, f)

[Lunatik00]

Thanks, they should add this to the export, to convert your results to onnx you need to have it in a pkl, I was unable to do it with a pth

[ilya1234-ai]

Thanks!