Pytorch加载模型时的常见问题

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

可能原因：

1. 加载模型时的网络与训练时的不一致（较常见）
2. 训练使用GPU，而测试时用的CPU
3. 在模型训练时使用了 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'])