我正在使用keras-rl通过D-DQN算法训练我的网络。我正在使用model.fit_generator()
函数在GPU上进行训练,以允许在执行反向传递时将数据发送到GPU。我怀疑与GPU处理数据的速度相比,数据的生成速度太慢。
按照D-DQN算法的指示,在生成数据时,我必须首先使用模型预测Q值,然后将这些值用于反向传播。如果使用GPU来运行这些预测,则意味着它们正在破坏我的数据流(我希望backprops尽可能频繁地运行)。
是否可以指定在哪个设备上运行特定操作的方法?这样我就可以在CPU上运行预测,而在GPU上运行backprops。
答案 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加载模型并进行预测了。