Tensorflow 2：如何在GPU和CPU之间切换执行？

Question

在使用独立的Keras 2.X和TensorFlow 1.X时，我曾经使用以下代码片段在GPU上进行训练，在CPU上运行推断（对于我的RNN模型来说，这种方式更快）：

keras.backend.clear_session()

def set_session(gpus: int = 0):
    num_cores = cpu_count()

    config = tf.ConfigProto(
        intra_op_parallelism_threads=num_cores,
        inter_op_parallelism_threads=num_cores,
        allow_soft_placement=True,
        device_count={"CPU": 1, "GPU": gpus},
    )

    session = tf.Session(config=config)
    k.set_session(session)

在 tensorflow 2.0 中，ConfigProto 功能不再可用（我正在使用集成的 tensorflow.keras）。在开始时，可以运行 tf.config.experimental.set_visible_devices() 以禁用 GPU，但任何后续对 set_visible_devices 的调用都会导致 RuntimeError：Visible devices cannot be modified after being initialized。是否有一种重新初始化可见设备或切换可用设备的方法？

- valend.in

Answer 1

你可以使用tf.device来明确设置所需的设备。例如：

import tensorflow as tf    

model = tf.keras.Model(...)

# Run training on GPU
with tf.device('/gpu:0'):
    model.fit(...)

# Run inference on CPU
with tf.device('/cpu:0'):
    model.predict(...)

如果您只有一个CPU和一个GPU，则上面使用的名称应该可以工作。否则，device_lib.list_local_devices() 可以给出您设备的列表。 这篇文章提供了一个很好的函数，用于列出设备的名称，我在此进行了修改以显示CPU：

from tensorflow.python.client import device_lib

def get_available_devices():
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos if x.device_type == 'GPU' or x.device_type == 'CPU']

- adamconkey

---

你知道TensorFlow默认是在GPU还是CPU上运行的吗？ - cloudscomputes

有没有一种方法可以为整个脚本的范围设置它？我的意思是，不使用with语句？类似于tf.set_device('/cpu:0')这样的东西？ - Lucas Azevedo

这可能是一个实际的解决方案，但请注意，一旦您输入 tf.device('/cpu:0'):，TF仍然会分配GPU内存。我认为这是TF的一个限制。 - Jongwook Choi

假设您只想在CPU上运行，您可以通过在初始化TensorFlow之前将环境变量CUDA_VISIBLE_DEVICES设置为空字符串来避免TensorFlow分配GPU内存的习惯。 - pygosceles

@LucasAzevedo 如果你想为整个范围设置设备，可以使用以下代码：physical_devices = tf.config.list_physical_devices('GPU'); tf.config.set_visible_devices(physical_devices[0], 'GPU')。这个功能从tensorflow 2.1版本开始实现。 - BeCurious

Answer 2

使用tf.device能够帮助你吗？

通过这种方式，你可以将一些操作设置在CPU或GPU上执行。

- AlexisBRENON

Answer 3

我只是重新启动内核，这对我起作用了

- Lowity