为啥 TF Keras 推理方式比 Numpy 操作慢？

Posted 2023-03-12

【中文标题】为啥 TF Keras 推理方式比 Numpy 操作慢？

【英文标题】：Why is TF Keras inference way slower than Numpy operations?为什么 TF Keras 推理方式比 Numpy 操作慢？

【发布时间】：2018-07-25 13:59:25

【问题描述】：

我正在研究使用 Keras 和 Tensorflow 实现的强化学习模型。我必须在单个输入上频繁调用 model.predict()。

在一个简单的预训练模型上测试推理时，我注意到使用 Keras 的 model.predict 比仅在存储的权重上使用 Numpy 慢得多。为什么这么慢，我该如何加速？对于复杂的模型，使用纯 Numpy 是不可行的。

import timeit
import numpy as np
from tensorflow.python.keras.models import Sequential
from tensorflow.python.keras.layers import Dense

w = np.array([[-1., 1., 0., 0.], [0., 0., -1., 1.]]).T
b = np.array([ 15., -15., -21., 21.])

model = Sequential()
model.add(Dense(4, input_dim=2, activation='linear'))
model.layers[0].set_weights([w.T, b])
model.compile(loss='mse', optimizer='adam')

state = np.array([-23.5, 17.8])

def predict_very_slow():
    return model.predict(state[np.newaxis])[0]

def predict_slow():
    ws = model.layers[0].get_weights()
    return np.matmul(ws[0].T, state) + ws[1]

def predict_fast():
    return np.matmul(w, state) + b

print(
    timeit.timeit(predict_very_slow, number=10000),
    timeit.timeit(predict_slow, number=10000),
    timeit.timeit(predict_fast, number=10000)
)
# 5.168972805004538 1.6963867129435828 0.021918574168087623
# 5.461319456664639 1.5491559107269515 0.021502970783442876

【问题讨论】：

您能提供您的模型定义吗？

一个有趣的见解。请提供有关您的训练/推理以及模型定义的更多信息（可能是更多代码？）。

如果我在一个10000行的数组上使用model.predict() 1次，大约需要0.1秒。

您使用的是 GPU 吗？

@MatiasValdenegro 不，我只使用 CPU。

【参考方案1】：

memory leak issue 似乎仍然存在于 Keras 中。该问题中提到的以下代码行对我有用：

import ... as K
import gc

model = ....
del model
K.clear_session()
gc.collect()

【参考方案2】：

有点晚了，但可能对某人有用：

将model.predict(X) 替换为model.predict(X, batch_size=len(X))

应该可以的。

【讨论】：

这节省了大量时间。唯一的问题是确保 batch_size 不是很大，否则 tensorflow 会抛出 OOM。

当其他一些更复杂的解决方案（例如使用已编译/未编译模型）不起作用时，这对我有用。

【参考方案3】：

您是否正在循环运行您的 Keras 模型（带有 TensorFlow 后端）？如果是这样，Keras 会在此处发现内存泄漏问题：LINK

在这种情况下，您必须导入以下内容：

import keras.backend.tensorflow_backend
import tensorflow as tf

from keras.backend import clear_session

最后，在完成计算后，您必须在循环的每次迭代结束时添加以下内容：

clear_session()
if keras.backend.tensorflow_backend._SESSION:
    tf.reset_default_graph()
    keras.backend.tensorflow_backend._SESSION.close()
    keras.backend.tensorflow_backend._SESSION = None

这应该可以帮助您在每个循环结束时释放内存，并最终加快处理速度。我希望这会有所帮助。

【讨论】：

据我了解，只有在循环中创建模型时，内存泄漏才是问题。我在测试用例中只使用了一个模型对象。

@schoeberl 明白了。可能有几种方法可以使单个模型更快，这取决于您如何设置模型。快速搜索让我找到了这个，我认为这可能会有所帮助：LINK