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

时间:2018-07-24 20:47:17

标签: python tensorflow keras keras-rl

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

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

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

2 个答案:

答案 0 :(得分: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

答案 1 :(得分:0)

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