Python Tensorflow 在jupyter中设置CUDA_VISIBLE_DEVICES

在本文中，我们将介绍如何在Python的TensorFlow库中设置CUDA_VISIBLE_DEVICES环境变量。CUDA_VISIBLE_DEVICES是一个用于设置可见的GPU设备的环境变量，用于指定哪些GPU设备将被TensorFlow库使用。这在使用多个GPU或在特定场景下只使用特定GPU时非常有用。

在使用TensorFlow进行深度学习任务时，我们经常需要使用GPU来加速计算。但是，在某些情况下，我们可能只想使用特定的GPU设备来运行我们的代码，而忽略其他的GPU设备。TensorFlow提供了一个环境变量CUDA_VISIBLE_DEVICES来控制代码可见的GPU设备。

阅读更多：Python 教程

查看可用的GPU设备

在设置CUDA_VISIBLE_DEVICES之前，我们首先需要了解系统上可用的GPU设备。我们可以使用TensorFlow的GPU设备API来查看系统上的可用GPU设备。

下面的示例代码演示了如何使用TensorFlow的tf.config.experimental.get_visible_devices函数来获取可见的GPU设备列表。

import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices('GPU')
visible_gpus = tf.config.experimental.get_visible_devices('GPU')

print("可见的GPU设备:")
for gpu in visible_gpus:
    print(gpu)

print("所有的GPU设备:")
for gpu in gpus:
    print(gpu)

运行上述代码，将输出可见的GPU设备列表和所有的GPU设备列表。可见的GPU设备列表是我们目前可以使用的GPU设备。

设置可见的GPU设备

设置CUDA_VISIBLE_DEVICES环境变量可以控制TensorFlow代码可见的GPU设备。可以使用以下两种方法来设置CUDA_VISIBLE_DEVICES环境变量。

方法一：设置环境变量

使用os库的environ函数可以设置环境变量。

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"  # 设置为设备0和设备1

# 在这里执行TensorFlow代码

上述代码将设置CUDA_VISIBLE_DEVICES环境变量为0,1。这将导致TensorFlow只看到设备0和设备1，并忽略其他的GPU设备。

方法二：使用tf.config.experimental.set_visible_devices函数

使用TensorFlow的tf.config.experimental.set_visible_devices函数也可以设置可见的GPU设备。

import tensorflow as tf

physical_devices = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_visible_devices(physical_devices[0:2], 'GPU')

# 在这里执行TensorFlow代码

上述代码将设置可见的GPU设备为系统中的设备0和设备1。只有这两个设备会被TensorFlow看到。

请注意，通过上述方法设置的CUDA_VISIBLE_DEVICES环境变量只对当前的Python进程有效。

示例

让我们通过一个示例来了解如何在jupyter中设置CUDA_VISIBLE_DEVICES。假设我们有两个GPU设备，我们只想在jupyter中使用其中一个设备。

首先，我们需要安装TensorFlow库，可以使用以下命令：

!pip install tensorflow

然后，我们可以在jupyter中运行以下代码来设置CUDA_VISIBLE_DEVICES环境变量并验证其效果。

import tensorflow as tf
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "1"  # 只使用设备1

# 检查可见的GPU设备
visible_gpus = tf.config.experimental.get_visible_devices('GPU')
print("可见的GPU设备:")
for gpu in visible_gpus:
    print(gpu)

# 在这里执行TensorFlow代码

上述代码将设备1设置为可见的GPU设备。然后，我们可以使用tf.config.experimental.get_visible_devices函数来检查可见的GPU设备，并打印出结果。

总结

通过设置CUDA_VISIBLE_DEVICES环境变量，我们可以在Python的TensorFlow库中控制可见的GPU设备。这在使用多个GPU或在特定场景下只使用特定GPU时非常有用。我们可以使用os.environ函数或TensorFlow的tf.config.experimental.set_visible_devices函数来设置可见的GPU设备。请注意，通过上述方法设置的CUDA_VISIBLE_DEVICES环境变量只对当前的Python进程有效。