Pytorch 模型load各种问题解决

最新推荐文章于 2023-04-17 11:08:08 发布

xddwz

最新推荐文章于 2023-04-17 11:08:08 发布

阅读量1k

点赞数

分类专栏： Pytorch 文章标签： python

本文链接：https://blog.csdn.net/xddwz/article/details/116206358

版权

Pytorch 专栏收录该内容

16 篇文章 2 订阅

订阅专栏

出现unexpected key module.xxx.weight问题

有时候你的模型保存时含有 nn.DataParallel时，就会发现所有的dict都会有 module的前缀。
这时候加载含有module前缀的模型时，可能会出错。其实你只要移除这些前缀即可

  pretrained_net = Net_OLD()
  pretrained_net_dict = torch.load(save_path)
  new_state_dict = OrderedDict()
  for k, v in pretrained_net_dict.items():
      name = k[7:] # remove `module.`
      new_state_dict[name] = v
  # load params
  pretrained_net.load_state_dict(new_state_dict)

总结
保存的Dict是按照net.属性.weight来存储的。如果这个属性是一个Sequential，我们可以类似这样net.seqConvs.0.weight来获得。
当然在定义的类中，拿到Sequential的某一层用[], 比如self.seqConvs[0].weight.
strict=False是没有那么智能，遵循有相同的key则赋值，否则直接丢弃。

参考：https://blog.csdn.net/Hungryof/article/details/81364487?utm_medium=distribute.pc_relevant.none-task-blog-2%7Edefault%7EBlogCommendFromMachineLearnPai2%7Edefault-1.control&dist_request_id=&depth_1-utm_source=distribute.pc_relevant.none-task-blog-2%7Edefault%7EBlogCommendFromMachineLearnPai2%7Edefault-1.control