Loading the pretrained weights into pytorch

[josh-wende]

Hi,

I'm trying to use the pretrained weights provided in the readme, but am having some trouble. I loaded model_final.pth as a state_dict, and then tried to load it into a new instance of KeypointDetector, but this hit an error because the expected keys for a KeypointDetector state_dict and the given dictionary keys do not match up. What am I doing wrong?

Thanks,
Josh

[gch]

It seems like the codebase may have changed since the pretrained weights were saved. I found the following (manually changed) key names gets things to work:

    pretrained = torch.load('/mnt/drive/MyDrive/Datasets/SMOKE/model_final.pth')['model']
    for k1, k2 in zip(sorted(pretrained.keys()), sorted(model.state_dict().keys())):
        assert(k1.replace('module.', '') == k2)
        model.state_dict()[k2][:] = pretrained[k1]

EDIT: reading the code further, there's just some extra information that's saved when the training script serializes the model (training debug data, etc.). It looks like

SMOKE/smoke/utils/model_serialization.py

Line 69 in bc5d2bb

def load_state_dict(model, loaded_state_dict):

provides an existing helper function that loads from the slightly customized serialization.