pytorch并行计算，显存溢出问题之torch并行问题解决办法

最新推荐文章于 2025-06-15 11:27:40 发布

kao_lengmian

最新推荐文章于 2025-06-15 11:27:40 发布

阅读量6.4k

点赞数 3

CC 4.0 BY-SA版权

分类专栏：分享文章标签： cuda gpu python

本文链接：https://blog.youkuaiyun.com/kao_lengmian/article/details/108492848

分享专栏收录该内容

9 篇文章

订阅专栏

本文解决了一个在使用PyTorch模型时遇到的显存超出问题，即使显存仍有大量空间。问题根源在于多GPU环境下，特定GPU内存不足导致模型加载失败。通过调整模型加载代码，指定使用具有足够显存的GPU，成功解决了问题。

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

使用的是torch模型，遇到一个显存超出报错问题，信息如下：

  File "/home/wangcf/anaconda3/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 443, in to
    return self._apply(convert)
  File "/home/wangcf/anaconda3/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 203, in _apply
    module._apply(fn)
  File "/home/wangcf/anaconda3/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 203, in _apply
    module._apply(fn)
  File "/home/wangcf/anaconda3/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 225, in _apply
    param_applied = fn(param)
  File "/home/wangcf/anaconda3/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 441, in convert
    return t.to(device, dtype if t.is_floating_point() else None, non_blocking)
RuntimeError: CUDA error: out of memory

但是此时显存可以查看到还有很多空间，明显不是显存不够，那是哪里出问题了？
显卡信息
我试过更换其他同类型模型，完全没有问题，可以跑通，但是换回来就又是显存溢出，那么肯定不是模型超大的问题了（模型确实不大），一度怀疑是模型损坏，但是经过检查，可以正常进行模型推理，那也不是损坏问题了。
这时脑子灵光一闪，是不是因为多显卡的问题，于是乎，改了一行code：

model.load_state_dict(torch.load(args.model_path, map_location='cuda:0'))

你没有看错，是显卡的分配问题，训练时使用的是4张显卡，而在使用的时候有一张显卡内存是不够用的，所以导致模型加载失败，那么只要让一张显卡加载模型即可。