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是pytorch使用多gpu训练时所使用的方法,但是使用nn.DataParallel之后,模型的读取就会有所不同。最常见的情况就是使用的预训练模型并不是在多张gpu上训练得来的,没有使用nn.DataParallel包装,但现在想要使用这个预训练模型就不能直接用下面的代码

model.load_state_dict(torch.load(save_path))

2.错误原因
使用nn.DataParallel包装后的模型参数的关键字会比没用nn.DataParallel包装的模型参数的关键字前面多一个“module.”,可以看看前面的错误,missing key 和 unexpected key的差别就在那个“module.”上。还有一种情况就是missing key 是module.features.0.weight,但unexpected key是 features.module.0.weight,就是module和feature的位置反过来了。
3.解决方法
(1)在使用nn.DataParallel之前就先读取模型,然后再使用nn.DataParallel,代码的顺序是

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

原文链接:https://blog.csdn.net/single_dog/article/details/114413417

你可能感兴趣的:(日常bug系列)