python - 在 python 中的运行脚本中关闭 GPU

Question

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

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

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

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

但是此代码不会在运行脚本中关闭 GPU。它只能在计算开始时使用。
任何想法都受到高度赞赏
提前致谢

Answer 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 设备，示例如下:

If you would like a particular operation to run on a device of your choice instead of what's automatically selected for you, you can use with tf.device to create a device context, and all the operations within that context will run on the same designated 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)

You will see that now a and b are assigned to CPU:0. Since a device was not explicitly specified for the MatMul operation, the TensorFlow runtime will choose one based on the operation and available devices (GPU:0 in this example) and automatically copy tensors between devices if required.

如果您想要 tf.matmul op 不使用 GPU，只需要将其移入前面 with 的范围内陈述。