'Profiler Report' - like PyTorch-Lighting in TF.Keras interface

TensorFlow General Discussion
, ,
1

I found the profiler report produced in the pytorch-lighting (pt) model training is convenient to debug training. AFAIK, tf has also a profiler which can be viewed via tensor board. In pt, it gives as follows:

Profiler Report

Action                  |  Mean duration (s)    |  Total time (s)
-----------------------------------------------------------------
on_epoch_start          |  5.993e-06            |  5.993e-06
get_train_batch         |  0.0087412            |  16.398
on_batch_start          |  5.0865e-06           |  0.0095372
model_forward           |  0.0017818            |  3.3408
model_backward          |  0.0018283            |  3.4282
on_after_backward       |  4.2862e-06           |  0.0080366
optimizer_step          |  0.0011072            |  2.0759
on_batch_end            |  4.5202e-06           |  0.0084753
on_epoch_end            |  3.919e-06            |  3.919e-06
on_train_end            |  5.449e-06            |  5.449e-06

Numbers are random of course. We can easily accomplish on_epoch_start/end, on_batch_start/end timing. But not sure about others (mostly). I don’t know yet if there is any such mechanism that exists already in tf. keras.

2

https://github.com/tensorflow/addons/issues/2427

4

@Bhack
Thanks for brining it up. Could you please give some suggestion on that?

5

I think that you can get the timing on all the available callbacks:

E.g.

6

@Bhack
Thank you for the link, helpful. As I also mentioned in the query that we can accomplish on_epoch_start/end, on_batch_start/end timing.

Actually, I was interested in timing on model_forward, model_backward , AFAIK, under normal circumstances, we can’t use callback for that, can we? (However, I will test other options using the callbacks to verify the timing results).

7

Currently we don’t have these callbacks:

But you could approximate with: