python - 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