Keras-如何在CPU上运行加载的模型

时间:2018-09-22 15:42:10

标签: python keras gpu cpu

我有带有在GPU上运行的tensorflow后端的keras。但是,我在训练LSTM,所以我在CPU上训练。

with tf.device('/cpu:0'):
    model = Sequential()
    model.add(Bidirectional(LSTM(50, return_sequences=True), input_shape=(50, len(train_x[0][0]))))
    model.add(TimeDistributed(Dense(1, activation='sigmoid')))
    model.compile(loss='binary_crossentropy', optimizer='Adam', metrics=['acc'])

我的问题是,当我保存并加载模型时,加载模型的预测函数执行得非常慢。经过一些定时测试后,我相信发生的事情是加载的模型在GPU而不是CPU上运行,因此速度很慢。我尝试在CPU上编译加载的模型,但这并不能加快速度:

model.save('test_model.h5')
new_model = load_model('test_model.h5')
with tf.device('/cpu:0'):
    new_model.compile(loss='binary_crossentropy', optimizer='Adam', metrics=['acc'])

是否有一种方法可以使加载的模型达到与新训练的模型相同的速度?新训练的模型快了将近五倍。感谢您的帮助。

1 个答案:

答案 0 :(得分:0)

使用要使用的设备加载模型:

with tf.device('/cpu:0'):
    new_model = load_model('test_model.h5')