RuntimeError: CUDA out of memory. Tried to allocate 96.00 MiB

最新推荐文章于 2024-07-11 16:36:54 发布

原创最新推荐文章于 2024-07-11 16:36:54 发布 · 2.7k 阅读

2 ·

本内容遵循CC 4.0 BY-SA版权协议

标签

#cuda #pytorch

收录于

当在拥有多个GPU的环境中进行深度学习时，可能会遇到CUDA内存不足的错误，即使显存显示为空闲。这可能是由于PyTorch与TensorFlow在同一GPU上运行冲突导致。解决方法是通过设置环境变量CUDA_VISIBLE_DEVICES指定使用特定的GPU。例如，使用2,3,6,7号GPU，可以设置os.environ['CUDA_VISIBLE_DEVICES']='2,3,6,7'，然后将模型放在DataParallel中，如model=nn.DataParallel(model, device_ids=[0,1])，确保使用的是空闲GPU。这样可以避免资源冲突并有效利用GPU资源。

使用多个GPU进行深度学习任务的时候有时会出现明明显卡内存还有空间，但是会提示RuntimeError: CUDA out of memory. Tried to allocate 96.00 MiB的情况。使用nvidia-smi查看显卡使用情况发现，空间足够。
参考了一些文章看总结得出可能原因：
pytorch与tensorflow如果在同一个GPU上运行会出现冲突，因此选择空闲的GPU即可。
具体设置为使用CUDA_VISIBLE_DEVICES限制一下使用的GPU.

# 使用八块GPU中的2，3，6，7号（从0开始）
os.environ['CUDA_VISIBLE_DEVICES'] = '2,3,6,7'
device = torch.device('cuda')

if torch.cuda.device_count() > 1:
    print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
    model = nn.DataParallel(model, device_ids=[0,1]) # 设置使用的GPU为0和1号，即上述可用GPU中的2和3号
model.to(device)

参考：