Add Lightning-like Profiler to Keras

[innat]

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

[innat]

cc. @bhack

[bhack]

We had a long thread at #552

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

[qlzh727]

The profiler solution on TF side is via the xprof tracing with TensorBoard, and keras has callbacks for it. Adding @rchao who is the owner of the callbacks.

https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras

[innat]

I've tried to use profiling plug-in from tensorboard before but tbh, I hardly felt productive to work with it. It's plug-in in tensorboard which need to install separately, but IMO keras should have more dedicated profiling tools.

If we look at the lightning-profiler, it gives several options (simple, advance, pytorch, xla or custom) to profile the code, which works great and gives easy to inspect the output and provide actionable feedback in quick.

[bhack]

What is your proposal? Do you want to have an alternative UI for the profiler portobufs?

https://github.com/tensorflow/tensorflow/tree/master/tensorflow/tsl/profiler/protobuf

[innat]

@bhack
Sorry, I don't have any proposal to present here. The thing I can request to have is from UX, described above. Looking back to here, it feels a complex issue in general. I am also interested to know what @rchao @fchollet thoughts about such possible feature.

[bhack]

Have you seen keras-team/keras-cv#141 (comment)?

[innat]

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

[bhack]

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

keras-team/keras-cv#17 (comment)

[innat]

cc. @rchao
A gentle reminder for this request.

[rchao]

This appears to be useful to me, but it's unclear whether there is available bandwidth on the team since it's some substantial work. Happy to check with the team and circle back in a couple of weeks.

[bhack]

@rchao Do you know something about the mentioned internal tool?

[rchao]

Correct, there has been experimental efforts on a profiling callback that hasn't been open sourced yet. That's one thing I'll also check with the team.