如何让Jupyter Notebook运行在多GPU中的某一指定GPU上？

Question

我在一台 Jupyter Notebook 中有一些 PyTorch 代码，需要在一个指定的 GPU（即不是“gpu 0”）上运行。因为GPU 0正在被其他人工作。我试过这个：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

但这不是工作！那么如何让代码在指定的GPU上成功运行呢？

Answer 1

您可以使用

CUDA_VISIBLE_DEVICES=1

启动 Jupyter Notebook。这通常是我在集群上执行的操作，因为 PyTorch 文档建议设置

CUDA_VISIBLE_DEVICES

与 torch

 等 

torch.cuda.set_device(device) 函数相比：

$ CUDA_VISIBLE_DEVICES=1 jupyter notebook &

您还可以使用

torch.cuda.is_available()

检查笔记本中可用的设备

Answer 2

您可以在包含将模型放置到设备上的代码的同一 Jupyter Notebook 单元的开头设置环境变量。 也可以改一下

不要写：

os.environ["CUDA_VISIBLE_DEVICES"]="4,5,6,7"

或者：

%set_env CUDA_VISIBLE_DEVICES=4,5,6,7

但是设置环境变量：

!export CUDA_VISIBLE_DEVICES=4,5,6,7

有多种方法可以将此设置设为默认设置：.bashrc、kernel.json、Jupyter 服务设置（全部未经测试），请参阅 Stack Overflow 如何在 Jupyter Notebook 中设置 env 变量。