在Jupyter Notebook中选择特定的GPU使用

Question

我正在一个Linux服务器上编写.ipynb文件。

我使用的Linux服务器有多个GPU，但我应该只使用空闲的GPU，以免意外中止其他人的程序。

我已经知道对于常见的.py文件，我们可以在命令行中添加一些指令来选择通用GPU（例如export CUDA_VISIBLE_DEVICES = #），但这对jupyter笔记本是否有效？如果无效，我该如何指定要使用的GPU。

- Arist12

Answer 1

你需要正确选择它的名称。例如，可能有三个可用的GPU设备，分别为"cuda:0"，"cuda:1"，"cuda:2"。要选择第三个设备，您需要运行以下代码：

if torch.cuda.is_available():
  dev = "cuda:2"
else:
  dev = "cpu"
device = torch.device(dev)

- sadegh arefizadeh

Answer 2

使用tensorflow如下

gpus = tf.config.list_physical_devices('GPU')
if len(gpus) > 0: 
   tf.config.experimental.set_visible_devices(gpus[0], 'GPU')

- noone