python - Tensorflow 模型预测很慢

Question

我有一个 TensorFlow单机型号Dense层:

model = tf.keras.Sequential([tf.keras.layers.Dense(2)])
model.build(input_shape=(None, None, 25))

我在 float32 中构建了一个输入向量:

np_vec = np.array(np.random.randn(1, 1, 25), dtype=np.float32)
vec = tf.cast(tf.convert_to_tensor(np_vec), dtype=tf.float32)

我想将其提供给我的模型进行预测，但速度非常慢。
如果我打电话 predict或 __call__与在 NumPy 中执行相同操作相比，它需要很长时间。

调用 %timeit model.predict(vec) :

10 loops, best of 3: 21.9 ms per loop

按原样调用模型 %timeit model(vec, training=False) :

1000 loops, best of 3: 806 µs per loop

自己执行乘法运算

weights = np.array(model.layers[0].get_weights()[0])   
%timeit np_vec @ weights

1000000 loops, best of 3: 1.27 µs per loop

使用火炬自己执行乘法

100000 loops, best of 3: 2.57 µs per loop

Google Colab: https://colab.research.google.com/drive/1RCnTM24RUI4VkykVtdRtRdUVEkAHdu4A?usp=sharing
如何使我的 TensorFlow 模型在推理时间内更快？
特别是因为我不仅有 Dense层，但我也使用 LSTM我不想在 NumPy 中重新实现它。

Answer 1

整个故事背后是在 Keras 中实现 LSTM 层。 Keras LSTM 层有一个默认参数 unroll=False .这会导致 LSTM 运行一个符号循环(循环会导致更多时间)。尝试向 LSTM 添加一个额外的参数作为 unroll=True .

tf.keras.layers.LSTM(64, return_sequences=True, stateful=True, unroll=True)

这可能会导致高达 2 倍的速度提升(在我的机器上测试，使用 %timeit model(vec, training=False))。但是，使用 unroll=True可能会导致为更大的序列占用更多的内存。更多查询，请查看 Keras LSTM documentation .