【Pytorch】加载模型避坑坑load_state_dict中的strict使用与加载多GPU训练的模型

背景

加载模型的小知识, 使用多GPU训练的模型并保存到ckpt中后,使用torch.load_state_dict加载模型的时候将会报错,但是如果将其中的参数strict设置为True的时候就可以加载,但是当使用加载后的模型去预测数据时,结果错的离谱。 相关内容可以看看这篇博文:关于Pytorch加载模型参数的避坑指南.

那么对应的解决方案是,在使用多GPU训练保存模型的时候,保存的模型应该是model.module,并不是直接保存model。(假设使用的模型是model)

你可能感兴趣的:(#,pytorch,pytorch,深度学习)