alexandrosstergiou / SoftPool

[ICCV 2021] Code for approximated exponential maximum pooling

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

How to load your pre-trained model?

caiduoduo12138 opened this issue · comments

when I load your pre-trained model, it seems to have some error.

"state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)"
"load_state_dict_from_url" does not define.

I download the model resnet-50_best.pth, and load it but error occurs.
Error(s) in loading state_dict for ResNet:
Missing key(s) in state_dict: "conv1.weight", "bn1.weight", "bn1.bias", "bn1.running_mean", "bn1.running_var", "layer1.0.conv1.weight", "layer1.0.bn1.weight", "layer1.0.bn1.bias", "layer1.0.bn1.running_mean", "layer1.0.bn1.running_var", ..........

Can give me some advice to load your pre-trained model. It would be nice if you could provide detailed code. Thanks.

Hi @caiduoduo12138 ,

Guessing that you also got : Unexpected key(s) in state_dict: "epoch", "arch", "state_dict", "best_acc1", "optimizer". as the last line in your error. Basically meaning that you should inlude the dictionary key state_dict during loading, e.g. :

net = torch.nn.DataParallel(resnet50(use_softpool=True))
net.load_state_dict(torch.load('weights/resnet-50_best.pth')['state_dict'])

Note that all models were trained with torch.nn.DataParallel, meaning that module names include a module wrapper. The example code above works regardless, but if you want to not use torch.nn.DataParallel, you could do something smarter like:

from collections import OrderedDict
n_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k.split('module.')[-1]
    new_state_dict[name] = v
model.load_state_dict(new_state_dict)

You can read more about data parallelism in torch here: [link]

Best,
Alex