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'])
Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐