在 python 的运行脚本中关闭 GPU

Posted 2023-03-23

【中文标题】在 python 的运行脚本中关闭 GPU

【英文标题】：Switch off the GPU in a running script in python

【发布时间】：2021-07-22 22:27:03

【问题描述】：

以下代码激活 GPU 0 并加载必要的库，如 Cuda 等。

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

当不再需要 GPU 时，如何在运行脚本中关闭它？例如，因为训练模型的评估需要在 CPU 上运行。 我试图改变环境变量。

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

但是这段代码并没有在运行脚本中关闭 GPU。只能在计算初期使用。

非常感谢任何想法

提前致谢

【问题讨论】：

这通常不会像您已经发现的那样工作。这种方式的器件选择是粗略的。它依赖于 CUDA 运行时在初始化时采样 一次 的环境变量。此后，对 env var 的更改无效。但是 CUDA 有其他方法来选择要使用的设备，即程序化的。因此，如果你想在 TF 中切换 GPU，一个更好的方法是让你的 TF 脚本be aware of multi-gpu usage（

这，我的朋友，实际上是我正在寻找的东西！谢谢

我也尝试使用-1。然后 list_physical_devices('GPU') 显示 0 并且我的 GPU 内存没有使用，但是....我的 GPU 处理器仍在使用，我的工作运行速度提高了 3 倍（小型神经网络）。

【参考方案1】：

在 python 中选择 GPU 的标准方法并不单一，这取决于您在 python 中使用的框架来访问 GPU。

您描述的environment variable 用法有些粗略，其行为与您描述的一样。环境变量会在 CUDA 初始化时影响 CUDA 行为，这通常只会发生一次，通常是在 python 脚本的开头。此后，对环境变量的更改将无效。 CUDA 在初始化时只对其采样一次。然而，在其他方面，这种效果应该或多或少是一致的，与您在 python 中使用 GPU 的方式无关。

除此之外，我们需要讨论一个特定的框架。底层 CUDA 技术 (CUDA C++) 允许框架开发人员动态控制可以在特定点使用哪些 GPU（来自那些通过环境变量“暴露”的 GPU；您不能覆盖它。）这种动态的方法暴露的控制会因框架而异。由于您的问题被标记为tensorflow，因此您可以找到关于 TF 如何公开 GPU 选择/控制here 的一个很好的概述。

具体来说，您要问的是不使用 GPU。我建议您参考“手动设备放置”section。请注意，此处的“设备”指的是 CPU 和 GPU，因此在此方法中“关闭”GPU 使用的方法是使用 CPU 设备，示例如下：

如果您希望特定操作在您选择的设备上运行，而不是自动为您选择的设备，您可以使用 tf.device 创建设备上下文，该上下文中的所有操作都将在相同的指定设备。

tf.debugging.set_log_device_placement(True)

# Place tensors on the CPU
with tf.device('/CPU:0'):
  a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

# Run on the GPU
c = tf.matmul(a, b)
print(c)

Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

您将看到现在 a 和 b 已分配给 CPU:0。由于没有为 MatMul 操作明确指定设备，TensorFlow 运行时将根据操作和可用设备（本示例中为 GPU:0）选择一个设备，并在需要时自动在设备之间复制张量。

如果您希望 tf.matmul 操作不使用 GPU，只需将其移动到前面的 with 语句的范围内即可。