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

只需使用 magics:

，无需任何导入即可更快地完成

%env CUDA_DEVICE_ORDER=PCI_BUS_ID
%env CUDA_VISIBLE_DEVICES=0

注意所有环境变量都是字符串，所以不需要使用"。您可以验证环境变量是否由运行: %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"

Tensorflow set CUDA_VISIBLE_DEVICES within jupyter