Tensorflow 在 jupyter 中设置 CUDA_VISIBLE_DEVICES

Posted 2023-02-23

【中文标题】Tensorflow 在 jupyter 中设置 CUDA_VISIBLE_DEVICES

【英文标题】：Tensorflow set CUDA_VISIBLE_DEVICES within jupyter

【发布时间】：2016-10-20 00:15:58

【问题描述】：

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

使用 CUDA_VISIBLE_DEVICES，我可以隐藏 python 文件的设备，但是我不确定如何在笔记本中执行此操作。

有没有办法将不同的 GPU 隐藏到运行在同一台服务器上的笔记本中？

【参考方案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

【讨论】：

对我来说它不起作用 os.environ["CUDA_VISIBLE_DEVICES"]="0" ...我将其更改为 os.environ["CUDA_VISIBLE_DEVICES"]="" 然后它起作用了。我正在使用张量流 1.12

我尝试过如前所述显示 tensorflow 本地设备。我的系统只输出设备类型：CPU。这是否意味着 tensorflow 没有运行 GPU。

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

这必须在你的 notebook 的开头完成。

【参考方案2】：

只需使用magics，您就可以在不导入任何内容的情况下更快地做到这一点：

%env CUDA_DEVICE_ORDER=PCI_BUS_ID
%env CUDA_VISIBLE_DEVICES=0

注意所有的环境变量都是字符串，所以不需要使用"。您可以通过运行：%env <name_of_var> 来验证是否设置了 env-variable。或者使用%env查看所有这些。

【讨论】：

这个env设置会影响其他python脚本吗？因为它看起来像改变 os 的环境变量。

@JenkinsY 您可以通过在一个脚本中设置变量并在另一个脚本中读取它来轻松检查它。

【参考方案3】：

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

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