Jupyter中的Tensorflow设置CUDA_VISIBLE_DEVICES


76

我有两个GPU,想同时通过ipynb运行两个不同的网络,但是第一个笔记本始终分配两个GPU。

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

无论如何,是否会将不同的GPU隐藏在同一服务器上运行的笔记本电脑中?

Answers:


156

您可以使用在笔记本中设置环境变量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

1
对我来说,它不起作用os.environ [“ CUDA_VISIBLE_DEVICES”] =“ 0” ...我将其更改为os.environ [“ CUDA_VISIBLE_DEVICES”] =“”然后它起作用了。我正在使用tensorflow 1.12
AlonSamuel

我尝试显示如上所述的tensorflow本地设备。我的系统仅输出设备类型:CPU。这是否意味着tensorflow没有运行GPU。
shaifali Gupta,

在导入tensorflow之后,这对我不起作用。
rjurney

33

您可以通过使用魔术来更快地完成此过程而无需任何导入:

%env CUDA_DEVICE_ORDER=PCI_BUS_ID
%env CUDA_VISIBLE_DEVICES=0

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


env设置会影响其他python脚本吗?就像改变os的环境变量一样。
JenkinsY

@JenkinsY您可以通过在一个脚本中设置变量并在另一个脚本中读取变量来轻松检查它。
萨尔瓦多·达利
By using our site, you acknowledge that you have read and understand our Cookie Policy and Privacy Policy.
Licensed under cc by-sa 3.0 with attribution required.