Which is the optimal way to measure the execution time of the layers on keras/tf when inferring?

Viewed 1144

I am trying to create an application that estimates the time consumed in the execution of each layer of a keras model. Although I get some time readings, they are not precise as they differ 500 times from each other. How can I retrieve the precise time required by each layer in the inference/predict step?

I have implemented a function to split the previous model in different models to be able to execute the model in different devices. Thus, I want to profile the application at a layer level.

I have tried two ways of retrieving the time consumed:

Increasing the verbosity of the call to predict function and the callback class you can see below.

273/273 [==============================] - 0s 13us/step (Verbosity=1)

Step time: 7.518984ms.(Callback)

I am running Python 3.7 and Keras 2.3 over TensorFlow 2.0. And I am not actually interested in performance, but in being capable of profiling models by layer.

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))

I feel like there is an easy and fancy way to do this but I am incapable of finding it...

Thanks!

0 Answers
Related