Missing key(s) in state_dict: “module.features.0.0.weight

Missing key(s) in state_dict: "module.features.0.weight", "module.features.0.bias", ....
Unexpected key(s) in state_dict: "features.0.weight", "features.0.bias", ...

nn.DataParallel is the method used by pytorch to train with multiple GPUs, but after using nn.DataParallel, the reading of the model will be different. The most common situation is that the pre-trained model used is not trained on multiple GPUs, and the nn.DataParallel package is not used, but if you want to use this pre-trained model, you cannot directly use the following code

model.load_state_dict(torch.load(save_path))

2. Reasons for the error
The keyword of the model parameter packaged with nn.DataParallel will be one more "module." before the keyword of the model parameter packaged with nn.DataParallel. You can look at the previous error, missing key and unexpected key The difference lies in the "module." Another situation is that the missing key is module.features.0.weight, but the unexpected key is features.module.0.weight, which means that the positions of the module and the feature are reversed.
3. Solution
(1) Read the model before using nn.DataParallel, and then use nn.DataParallel, the order of the code is

model.load_state_dict(torch.load(save_path))
model = nn.DataParallel(model, device_ids=[0, 1]) 

Original link: https://blog.csdn.net/single_dog/article/details/114413417

Guess you like

Origin blog.csdn.net/qq_35037684/article/details/115109416