load pretrain weight failed

[wilbur-caper]

import torch
import models

model = models.MobileViT_S()
PATH = "./MobileVit-S.pth.tar"
weights = torch.load(PATH, map_location=lambda storage, loc: storage)
model.load_state_dict(weights['state_dict'])
model.eval()
torch.save(model, './model.pt')

• I try to load the pre-train weight to test one demo; but the network structure does not seem to match the weights, is there any solution?

[wilile26811249]

Hi, @hererookie sorry for reply so late. Because, I wrapped this model into DataParallel when training, which stores the model in module.

Sol: model.load_state_dict({k.replace('module.',''):v for k,v in weights['state_dict'].items()}

Ref: https://discuss.pytorch.org/t/solved-keyerror-unexpected-key-module-encoder-embedding-weight-in-state-dict/1686/3

[wilbur-caper]

Hi, @hererookie sorry for reply so late. Because, I wrapped this model into DataParallel when training, which stores the model in module.

Sol: model.load_state_dict({k.replace('module.',''):v for k,v in weights['state_dict'].items()}

Ref: https://discuss.pytorch.org/t/solved-keyerror-unexpected-key-module-encoder-embedding-weight-in-state-dict/1686/3

thanks for your reply, I will try