深度学习中指定特定的GPU使用

news2024/9/24 9:22:36

前言

老生常谈，同样的问题，主要来源于：RuntimeError: CUDA error: out of memory

当使用完之后，想从其他方式调试，具体可看我这篇文章的：出现 CUDA out of memory 的解决方法

如果还有的空间使用，可看我下面的方式

1. 问题所示

执行代码的时候，出现如下问题：

  File "/home/l228/anaconda3/envs/stegan3.7/lib/python3.7/site-packages/torch/nn/modules/module.py", line 381, in to
    return self._apply(convert)
  File "/home/l228/anaconda3/envs/stegan3.7/lib/python3.7/site-packages/torch/nn/modules/module.py", line 187, in _apply
    module._apply(fn)
  File "/home/l228/anaconda3/envs/stegan3.7/lib/python3.7/site-packages/torch/nn/modules/module.py", line 187, in _apply
    module._apply(fn)
  File "/home/l228/anaconda3/envs/stegan3.7/lib/python3.7/site-packages/torch/nn/modules/module.py", line 193, in _apply
    param.data = fn(param.data)
  File "/home/l228/anaconda3/envs/stegan3.7/lib/python3.7/site-packages/torch/nn/modules/module.py", line 379, in convert
    return t.to(device, dtype if t.is_floating_point() else None, non_blocking)
RuntimeError: CUDA error: out of memory

截图如下所示：

在这里插入图片描述

最终查看是否还有已存的显卡供使用的时候，发现第一个显卡占用满了，但是第二个显卡不会自动使用：nvidia-smi

在这里插入图片描述

2. 解决方法

通过执行torch.cuda.get_device_name(id) 也可看到显卡的各个信息。单机多显卡：

在这里插入图片描述

方案一：通过命令行（最简单的方式）

使用命令如下：

CUDA_VISIBLE_DEVICES=1 python main.py

截图如下所示：

在这里插入图片描述

方案二：通过代码

指定多个空闲的GPU，供选择：

os.environ['CUDA_VISIBLE_DEVICES'] = "0, 1"
model = torch.nn.DataParallel(model, device_ids=[0, 1]) 
device = torch.device("cuda:0" if torch.cuda.is_available() and not args.no_cuda else "cpu")  # cuda 指定使用GPU设备

指定特定空闲的GPU，供选择：

os.environ['CUDA_VISIBLE_DEVICES'] = "1"
model = torch.nn.DataParallel(model, device_ids=[1])  
device = torch.device("cuda:1" if torch.cuda.is_available() and not args.no_cuda else "cpu")  # cuda 指定使用GPU设备