Jupyter中的Tensorflow设置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