Add Lightning-like Profiler to Keras

System information.

TensorFlow version (you are using): 2.9 Are you willing to contribute it (Yes/No) : No

Describe the feature and the current behavior/state

In PyTorch-Lightning, you can identify bottlenecks in your code using Profiler API. And I’ve found it super useful and easy to inspect the pros and cons of the code. This API offers 3 type of inferface, namely simple, advance and pytorch. I mostly use simple that summarize the execution time of medium to high level training + validation + inference API. The saved log files look something as follows.

# by passing a string
trainer = Trainer(..., profiler="simple")

# or by passing an instance
from pytorch_lightning.profiler import SimpleProfiler

profiler = SimpleProfiler()
trainer = Trainer(..., profiler=profiler)

I don’t know, in keras, if we have anything like this already. By having such high-level functionality in keras would really great. It will help to inspect the bottleneck of dataloader, model, callbacks etc, by reporting the time and memory consumption in the system.

I looked for it before, discussion.

Will this change the current api? How?

This can be considered as a Callback, I think.

tf.keras.callbacks.Profiler
or
tf.keras.utils.Profiler

Who will benefit from this feature?

The keras community. Engineers and researchers, both. Mostly engineers in production phase for code optimization.

Contributing

• Do you want to contribute a PR? (yes/no):
• If yes, please read this page for instructions
• Briefly describe your candidate solution(if contributing):

[bhack]

We had a long thread at https://github.com/keras-team/keras/issues/16699

Also using the standard tools it is still hard to precisely profile single tf.function: https://github.com/tensorflow/profiler/issues/503

[bhack]

No, haven’t seen. Is it generated from tf profiler? Or, custom from?

https://github.com/keras-team/keras-cv/issues/17#issuecomment-1021962006