pytorch加载分布式方式训练的模型

huahuahuahhhh

已于 2023-10-30 13:37:15 修改

阅读量1k

点赞数 1

CC 4.0 BY-SA版权

分类专栏： # Pytorch 文章标签： pytorch 分布式人工智能

于 2022-09-24 18:02:32 首次发布

本文链接：https://blog.youkuaiyun.com/Blankit1/article/details/127028571

Pytorch 专栏收录该内容

14 篇文章

订阅专栏

当使用PyTorch的load_state_dict加载分布式训练的模型时，若`strict=False`模型能加载但输出可能不正确。若`strict=True`，需移除'模块.'前缀以匹配权重。解决方案是创建新字典，将旧键转换。在加载后，模型预测结果与原始模型不符，可能原因包括未调用`model.eval()`、模型结构中存在dropout等。参考链接提供了详细排查步骤。

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

出现的问题

使用load_state_dict函数，如果strict=False,模型可以加载，但是模型的输出不对
如果strict=True,但是没有新建字典，模型无法加载

问题解决

保存的模型使用分布式方式训练，如果加载后模型不用分布式，则需要修改模型的key.
分布式模型的权值名字前面有module.，非分布式不包含。
比如分布式模型的权值名称名称为module.blocks.0.norm1.weight ，非分布式模型权值名称名称blocks.0.norm1.weight
需要重新建一个字典，用去掉module.的名字座位key,值不变

代码示例

chk_filename = os.path.join(args.checkpoint, args.resume if args.resume else args.evaluate)
print('Loading checkpoint', chk_filename)
checkpoint = torch.load(chk_filename, map_location=lambda storage, loc: storage)
checkpoint = checkpoint['model_pos']
new_checkpoint = {} ## 新建一个字典来访模型的权值
# print(checkpoint)
for k,value in checkpoint.items():
    key = k.split('module.')[-1]
    new_checkpoint[key] = value
    # print(k,key)

# model_pos.load_state_dict(checkpoint['model_pos'], strict=True)
model_pos.load_state_dict(new_checkpoint, strict=True)