Tensorflow在jupyter

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变量都是字符串，因此无需使用"。您可以通过运行%env <name_of_var>来验证是否已设置env变量。或者使用%env检查所有内容。

Answer 3

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

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