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

In tensorflow1.X 独立keras2.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)

This ConfigProto功能不再可用tensorflow2.0（我正在使用集成的tensorflow.keras）。一开始，可以运行tf.config.experimental.set_visible_devices()为了例如禁用 GPU，但任何后续调用set_visible_devices导致RuntimeError: Visible devices cannot be modified after being initialized。有没有一种方法可以重新初始化可见设备，或者是否有另一种方法可以切换可用的设备？

您可以使用tf.device https://www.tensorflow.org/api_docs/python/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() https://github.com/tensorflow/tensorflow/blob/d42facc3cc9611f0c9722c81551a7404a0bd3f6b/tensorflow/python/client/device_lib.py#L27可以为您提供设备列表。这个帖子 https://stackoverflow.com/questions/38559755/how-to-get-current-available-gpus-in-tensorflow提供了一个很好的函数来仅列出名称，我在这里对其进行了调整以显示 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']