Issues when loading MobileViT_S()

[Sehaba95]

I wanted to load the MobileViT_S() model and use the pre-trained weights, but I have got some errors in my code. To make it easier and help others, I will share my solution (in case there will be someone who is beginner like me):

def load_mobilevit_weights(model_path):
  # Create an instance of the MobileViT model
  net = MobileViT_S()
  
  # Load the PyTorch state_dict
  state_dict = torch.load(model_path, map_location=torch.device('cpu'))['state_dict']
  
  # Since there is a problem in the names of layers, we will change the keys to meet the MobileViT model architecture
  for key in list(state_dict.keys()):
    state_dict[key.replace('module.', '')] = state_dict.pop(key)
  
  # Once the keys are fixed, we can modify the parameters of MobileViT
  net.load_state_dict(state_dict)
  
  return net

net = load_mobilevit_weights("MobileViT_S_model_best.pth.tar")

[wilile26811249]

Thanks for your feedback.
May I share your code in the ReadMe?
It's very detail.
I will update in the readme.

[Sehaba95]

Hi, Yes sure, it is my pleasure. Thanks for your response and thank you for taking my request into consideration. Best Regards, Mohammed Sehaba

…

On Tue, 26 Apr 2022 at 10:30, Hong-Jia Chen ***@***.***> wrote: Hi, @Sehaba95 <https://github.com/Sehaba95> . May I share your code in the ReadMe? It's very detail. — Reply to this email directly, view it on GitHub <#7 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AETQY4VVOX7EDK4YBHIGJTLVG6STHANCNFSM5T3MLQHA> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[wilile26811249]

Hi @Sehaba95 , I have already updated your code in the ReadMe.
You can check it.
If you have any other questions, plz let me know. thank you.

[Sehaba95]

Thank you for adding my solution in your Readme file.
If I will encounter any problem, I will get back to you, and discuss it together.