DeBug Record——Pytorch设置默认CUDA不是0号

Pytorch设置默认CUDA不是0号

重新整理

返回来看自己之前的这篇文章，感觉写的逻辑太混乱了，回来简化一下。原先的版本留在文章结尾。

问题

Pytorch的版本是1.2.0，不知道现在改了没。
我希望能够只调用3号GPU，但是在用Pytorch库调用GPU时，使用以下代码，都会占用一些0号GPU显存：

# 版本一

with torch.cuda.device(3):
    main(opt)

# 版本二 设置环境变量（并没起作用，当时Pytorch版本为1.2.0）
import torch
import os
os.environ['CUDA_VISIBLE_DEVICES']='3'

解决方案

两种方案：
CUDA_VISIBLE_DEVICES=3 python script.py
或者

import os
os.environ['CUDA_VISIBLE_DEVICES']='3'
# 在torch之前设置好环境变量
import torch

原先的版本

问题说明

实验室服务器有4个GPU，同学用了0号，0号显存基本占满，剩余3个GPU，在代码里只使用了CUDA:3，但是cuda总是说memory不足。
报错：

cuda runtime error (77): an illegal memory access was encountered

或者

cuda runtime error (2) : out of memory

原因

在进行一行一行地debug时，发现，在运行以下代码时是没有问题的

device = torch.device('cuda:3' if torch.cuda.is_available() or 'cpu')
...
...
net = Net()
net.to(device)
...
...
input = input.to(device)

当时我就蒙了，这模型上cuda了，数据也上cuda了，这还没报错？那到底还有哪里能出错？
仔细一看，居然是下面这一行代码导致出错：

output = net(input)

经过一些资料的搜集，发现说是初始化时会默认占用0号卡…虽然我现在这份代码好像已经不算初始化的范畴了，不过还是死马当活马医，最后还真成了…

解决方案

参考了一些经验：
不管怎么指定pytorch都使用gpu0进行训练解决方法

为什么Pytorch死活要用第0块显卡–我和pytorch的恩怨情仇
具体方案：
在运行内容前面加上with torch.cuda.device(x):，如：

if __name__ == '__main__':
    opt = opts().parse()
    with torch.cuda.device(3):
        main(opt)

一些深入的小实验

首先来看这部分：

>>> with torch.cuda.device(3):
...     torch.cuda.device_count()
... 
4

这意味着，即使设置了torch.cuda.device(3)，程序还是能够看到别的显卡的。（对应前面参考的文章中，另一种方法是设置环境变量CUDA_VISIBLE_DEVICES，理论上使得看不到别的显卡，虽然我也试了这个环境变量设为3，不过能见到的显卡还是4个而不是1个。详情参见结尾）
接下来：

>>> a = torch.tensor([1,2,3])
>>> a.is_cuda
False
>>> with torch.cuda.device(3):
...     device = torch.device('cuda')
...     b = a.to(device)
... 
>>> b.is_cuda
True

在这里，经过查看，b占用的是3号卡的显存，说明torch.cuda.device可以将默认的cuda指定为特定的卡。
（注：我是用的Ubuntu，另开一个命令行输入命令nvidia-smi可以看到显存占用情况）

>>> a.is_cuda
False
>>> with torch.cuda.device(3):
...     device = torch.device('cuda:1')
...     c = a.to(device)
... 
>>> c.is_cuda
True

然后，上面运行后，c是在1号卡上的，说明在使用了torch.cuda.device应当是不影响别的卡的分配的。

一些遗留问题

在尝试的时候还遇到过另一种问题，就是比如说调用2、3号GPU时还是会出错，只调用3号GPU倒是没问题…
因为是调的别人的代码，也实在是懒得纠结了…最后乖乖只用3号GPU了…

结

祝各位DeBug顺利~

补

在深入小实验那一小节提到了环境变量CUDA_VISIBLE_DEVICES，当时我直接使用代码

import torch
...
...
import os
os.environ['CUDA_VISIBLE_DEVICES']='3'
print(torch.cuda.device_count())

显示的结果是4，说明CUDA_VISIBLE_DEVICES设置没起作用。
最后发现，CUDA_VISIBLE_DEVICES必须在import torch之前就进行设置，无论是命令行中

CUDA_VISIBLE_DEVICES=3 python script.py

还是

import os
os.environ['CUDA_VISIBLE_DEVICES']='3'
import torch
...

都可以