python指定gpu运行_PyTorch：从指定GPU开始运行

PyTorch默认使用从0开始的GPU，今天在Ubuntu上跑程序时出现了一个错误，如下：RuntimeError: CUDA out of memory. Tried to allocate 16.00 MiB (GPU 0; 1.95 GiB total capacity; 1.07 GiB already allocated; 9.25 MiB free; 16.95 MiB cached)

看了一下：

watch -n 0.1 nvidia-smi

原因是在我的Ubuntu系统中，0号GPU是Quadro P620(1998MB)，主要负责图形界面，1-8号是Telsa K80(11441MB)，主要负责计算，因此一旦真要跑个程序就会崩。(但在Windows里Quadro P620是8号GPU，/捂脸。)

解决方案如下：在终端中设定：

CUDA_VISIBLE_DEVICES=1,2,3,4,5,6,7,8 python my_script.py

2. 在python代码中设定：

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3,4,5,6,7,8"

#将模型放入GPU中

if torch.cuda.device_count() > 1:

model = torch.nn.DataParallel(model, device_ids=[0,1,2,3,4,5,6,7])

3. 在python代码中设定：

if torch.cuda.device_count() > 1:

model = torch.nn.DataParallel(model, device_ids=[1,2,3,4,5,6,7,8])

第2种和第3种方法等效。