RuntimeError: storage has wrong size: expected XXX got XXX 以及多卡GPU训练转到多卡GPU测试

程序员毛师傅

于 2022-04-15 14:03:20 发布

阅读量3.8k

点赞数

分类专栏：深度学习文章标签：算法计算机视觉机器学习

本文链接：https://blog.youkuaiyun.com/qq_42035274/article/details/124193207

版权

在进行深度学习模型训练时遇到RuntimeError，原因是DistributedDataParallel的模型权重保存不当。解决方法是在训练时只保存local_rank=0的权重。在从8张GPU训练模型转移到4张或单张GPU测试时，需通过调整torch.load的map_location参数适应不同的设备配置。

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

在训练时，我的代码出现了以下问题：

Traceback (most recent call last):
  File "/root/PycharmProjects/test.py", line 8, in <module>
    model_dict = torch.load("0.pth")
  File "torch/serialization.py", line 529,