PFLD_68points_Pytorch
PFLD_68points_Pytorch copied to clipboard
torch.nn.modules.module.ModuleAttributeError: 'MobileNetV2' object has no attribute 'module'
This happens when try to use test.py to load the model:
pfld_backbone.load_state_dict(torch.load(args.model_path, map_location= device))
To fix this issue, we have to change save model in the "train_model.py"
save model
checkpoint_path = os.path.join(model_dir, 'model_'+str(epoch)+'.pth') if args.all_model: torch.save(model, checkpoint_path)
as torch.save(model.state_dict(), checkpoint_path)
after that all works.
I am also wondering about an alternative without changing the train_model.py. So instead of using the train_model.py, we could do the code like below:
model = torch.load(pretrained_model) test(test_loader, model, args, device) ...