具有Tensorflow后端的Keras - 运行在CPU上预测但适合GPU

Posted 2021-04-05

我使用keras-rl用D-DQN算法训练我的网络。我正在使用model.fit_generator()函数在GPU上运行我的训练，以便在执行backprop时将数据发送到GPU。我怀疑与GPU处理数据的速度相比，数据生成速度太慢。

在生成数据时，按照D-DQN算法的指示，我必须首先用我的模型预测Q值，然后使用这些值进行反向传播。如果GPU用于运行这些预测，则意味着它们正在打破我的数据流（我希望backprops尽可能频繁地运行）。

有没有办法指定在哪个设备上运行特定的操作？在某种程度上，我可以运行CPU上的预测和GPU上的backprops。

答案

如果没有看到您的代码，很难正确回答您的问题。

下面的代码显示了如何列出可用设备和强制张量流以使用特定设备。

def get_available_devices():
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos]

get_available_devices()

with tf.device('/gpu:0'):
    //Do GPU stuff here
with tf.device('/cpu:0'):
    //Do CPU stuff here

另一答案

也许你可以在训练结束时保存模型。然后启动另一个python文件并在导入任何keras或tensorflow之前编写os.environ["CUDA_VISIBLE_DEVICES"]="-1"be。现在您应该能够加载模型并使用CPU进行预测。