python - 推断时测量 keras/tf 上的层执行时间的最佳方法是什么？

Question

我正在尝试创建一个应用程序来估计执行 keras 模型的每一层所消耗的时间。虽然我得到了一些时间读数，但它们并不精确，因为它们彼此相差 500 倍。如何检索推理/预测步骤中每一层所需的精确时间？

我已经实现了一个功能，将以前的模型拆分为不同的模型，以便能够在不同的设备中执行模型。因此，我想在层级分析应用程序。

我尝试了两种检索时间消耗的方法：

增加调用预测函数和回调类的详细程度，您可以在下面看到。

273/273 [===============================] - 0s 13us/步（详细度=1）

步长：7.518984ms。（回调）

我在 TensorFlow 2.0 上运行 Python 3.7 和 Keras 2.3。而且我实际上对性能并不感兴趣，而是对能够逐层分析模型感兴趣。

pred = model.predict(next_step, batch_size=256, verbose=1, callbacks=[MyCustomCallback()])

class MyCustomCallback(tf.keras.callbacks.Callback):

   def on_predict_begin(self, batch, logs=None):
        set_batch_time(time.time_ns())

   def on_predict_end(self, batch, logs=None):
        curr_time = time.time_ns()

        diff_time = curr_time - get_batch_time()
        print('Predict {} duration: {}'.format(self.model.layers[1].name, diff_time))

我觉得有一种简单而奇特的方法可以做到这一点，但我无法找到它......

谢谢！

标签： pythonpython-3.xtensorflowkerasprofiling