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

Question

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

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

是否可以指定在哪个设备上运行特定操作的方法？这样我就可以在CPU上运行预测，而在GPU上运行backprops。

Answer 1

很难在不看到代码的情况下正确回答您的问题。

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

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

Answer 2

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