解决pytorch中DataParallel后模型参数出现问题的方法

最新推荐文章于 2025-06-13 15:37:31 发布

转载最新推荐文章于 2025-06-13 15:37:31 发布 · 1w 阅读

Pytorch学习专栏收录该内容

46 篇文章

订阅专栏

本文介绍了解决PyTorch中使用DataParallel训练的模型在单GPU环境下加载时出现的'KeyError'问题的方法。通过创建一个新的OrderedDict，并移除模型state_dict中'模块.'前缀，使模型能够正确加载。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

pytorch中如果使用DataParallel，那么保存的模型key值前面会多处’modules.’，这样如果训练的时候使用的是多GPU，而测试的时候使用的是单GPU，模型载入就会出现问题。
一个解决方法是测试的时候强行DataParallel，但是有时候情况较为复杂，可以使用如下的方法：
（参考来源：https://discuss.pytorch.org/t/solved-keyerror-unexpected-key-module-encoder-embedding-weight-in-state-dict/1686/13）

original saved file with DataParallel

state_dict = torch.load(‘myfile.pth.tar’)

create new OrderedDict that does not contain module.

from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove module.
new_state_dict[name] = v
load params

1
model.load_state_dict(new_state_dict)
简而言之，就是重新创建一个OrderedDict，然后将它载入模型就行了。