Question

在带有独立tensorflow 2.X的keras 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)

此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。有没有一种重新初始化可见设备的方法，还是有另一种切换可用设备的方法？

Answer 1

使用tf.device可以帮助您吗？

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

Answer 2

您可以使用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']

Answer 3

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

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

3 个答案: