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)

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

1. 电话%timeit model.predict(vec)：

   10 个循环，最好的 3 个：每个循环 21.9 毫秒

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

   1000 个循环，最好的 3 个：每个循环 806 µs

3. 自己执行乘法运算

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

   1000000 个循环，最好的 3 个：每个循环 1.27 µs

4. 使用火炬自己执行乘法

   100000 个循环，最好的 3 个：每个循环 2.57 µs

谷歌 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文档。