Pytorch加载模型时的常见问题

伟大的兔兔

已于 2022-03-11 10:23:27 修改

阅读量3.7k

点赞数 1

分类专栏：科研技巧文章标签： pytorch 深度学习 python 模型加载

于 2022-03-10 10:21:07 首次发布

本文链接：https://blog.csdn.net/caijingxiang/article/details/123394688

版权

科研技巧专栏收录该内容

4 篇文章 0 订阅

订阅专栏

当使用Pytorch加载模型时出现'Unexpected keys in state_dict'错误，通常是因为模型结构不匹配、GPU与CPU环境差异或DataParallel处理。解决方法包括检查模型参数一致性，加载时指定map_location为'cpu'，以及对于DataParallel模型，加载state_dict内的'state_dict'子项。确保在加载和训练时模型环境的一致性。

摘要由CSDN通过智能技术生成

Pytorch加载模型报错 RuntimeError: Error(s) in loading state_dict for xxx: Unexpected key(s) in state_dict: xxx

问题描述：

RuntimeError: Error(s) in loading state_dict for xxx:  
Unexpected key(s) in state_dict: xxx

可能原因：

加载模型时的网络与训练时的不一致（较常见）
训练使用GPU，而测试时用的CPU
在模型训练时使用了 model = torch.nn.DataParallel(model) 并行训练

问题解决：

情况一：

这种情况没啥好说的，自己仔细检查检查，看看网络参数是否一致。

情况二：

# 加载模型时加上参数map_location='cpu'即可
# 如：

# 创建模型
model = TAANet_base(enc_dim=256, feature_dim=64, hidden_dim=128, layer=8, segment_size=200)
# 加载训练好的参数
model_state = torch.load('saves/temp/temp_best.pth.tar', map_location='cpu')
model.load_state_dict(model_state)

情况三：

# 创建模型
model = TAANet_base(enc_dim=256, feature_dim=64, hidden_dim=128, layer=8, segment_size=200)

model = torch.nn.DataParallel(model)

# 加载训练好的参数
model_state = torch.load('saves/temp/temp_best.pth.tar')
model.load_state_dict(model_state['state_dict'])