Tensorflow 2：如何将执行从 GPU 切换到 CPU 并返回？

Posted 2023-02-16

【中文标题】Tensorflow 2：如何将执行从 GPU 切换到 CPU 并返回？

【英文标题】：Tensorflow 2: how to switch execution from GPU to CPU and back?

【发布时间】：2020-04-12 03:59:41

【问题描述】：

在tensorflow 1.X 和独立的keras 2.X 中，我曾经使用以下 sn-p 在 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)

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

【参考方案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() 可以为您提供设备列表。 This post 提供了一个很好的函数，可以只列出名称，我在这里调整它也显示 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']

【讨论】：

你知道tensorflow默认是运行在GPU还是CPU上？

有没有办法为脚本的整个范围设置它？我的意思是，没有 with 语句？ tf.set_device('/cpu:0') 之类的东西？

【参考方案2】：

我只是重新启动内核，这对我有用

【参考方案3】：

使用tf.device 对您有帮助吗？

这样，您可以在 CPU 或 GPU 上设置一些操作。