Tensorflow在jupyter中设置CUDA_VISIBLE_DEVICES

Question

我有两个GPU,并希望同时通过ipynb运行两个不同的网络,但第一个笔记本总是分配两个GPU.

使用CUDA_VISIBLE_DEVICES,我可以隐藏python文件的设备,但是我不确定如何在笔记本中这样做.

反正是否有不同的GPU隐藏在同一台服务器上运行的笔记本电脑？

Answer 1

您可以使用在笔记本中设置环境变量os.environ.在初始化TensorFlow之前执行以下操作以将TensorFlow限制为第一个GPU.

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"]="0"

您可以仔细检查TF是否可以看到正确的设备

from tensorflow.python.client import device_lib
print device_lib.list_local_devices()

我倾向于使用像notebook_util这样的实用程序模块

import notebook_util
notebook_util.pick_gpu_lowest_memory()
import tensorflow as tf

Answer 2

只需使用魔法,你就可以更快地完成它而无需任何进口:

%env CUDA_DEVICE_ORDER=PCI_BUS_ID
%env CUDA_VISIBLE_DEVICES=0

请注意,所有env变量都是字符串,因此无需使用".您可以验证ENV变量是通过运行设置:%env <name_of_var>.或者检查所有这些%env.

Answer 3

您还可以启用多个 GPU 内核，如下所示：

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="0,2,3,4"