Profile particular functions / lines
shradhasehgal opened this issue · comments
Hey, is there a way to profile particular functions or code lines with one profiler i.e. not to have separate with profile as..statements around each of them?
Something similar to the NVIDIA nvtx markers.
Use case:
Want to profile only particular activity such as optimizer.step() or loss.backward() in a training loop, and not the entire loop.
@chaekit @aaronenyeshi Does Kineto have a profiler.start/stop API?
@anupambhatnagar Yes it is possible both in Python and C++, but it is not a public API.
You can manually get it to work using the underlying APIs.
Here is the code pointer to the underlying Python API:
https://github.com/pytorch/pytorch/blob/main/torch/autograd/profiler.py#L263
For Python:
- To start tracing:
- _prepare_profiler
- _enable_profiler
- To end tracing:
- torch.cuda.synchronize()
- results = _disable_profiler()
- To export results:
- results.save(your_local_path)
Here is the code pointer to the underlying C++ API:
https://github.com/pytorch/pytorch/blob/main/torch/csrc/autograd/profiler_kineto.h
This is also possible in C++:
- using namespace torch::autograd::profiler;
- To start tracing:
- prepareProfiler
- enableProfiler
- To end tracing:
- set a barrier
- auto result = disableProfiler();
- To export results:
- result->save("your_local_path");
@shradhasehgal please see above. i'm going to close the issue for now. feel free to open it if you have more questions related to this.
@anupambhatnagar Thank you! I will try this.