Adding a memory profiling example #41
Conversation
sanketpurandare
changed the title
|
Hey! |
dispatch_mem_profiler.py
Outdated
| def reduce_to_scalar_loss(inp): | ||
| return inp.sum() | ||
|
|
||
| from functorch.compile import aot_module, nop, print_compile, min_cut_rematerialization_partition |
There was a problem hiding this comment.
Can you move this to the top? Also try running black formatter just for consistent styling.
There was a problem hiding this comment.
The lint should already check that everything is black compliant:
Line 3 in 4214538
Since it passes on the PR, I guess it is good.
dispatch_mem_profiler.py
Outdated
|
|
||
| def __torch_dispatch__(self, func, types, args=..., kwargs=None): | ||
|
|
||
| global mem_usage, operator_names |
There was a problem hiding this comment.
You don't need this since you only write into these dictionaries.
Same for the other functions below except the ones that do set mem_usage.
There was a problem hiding this comment.
Cleaned up global usage
|
Could you add the new requirements so that we can merge this? |
Added them |
|
Python builtin Modules don't need to be in there ;) |
|
@albanD Can we merge this PR if it looks good? |
This example shows how to use torch_dispatch for memory profiling
It shows two ways
It also provides utilities to plot the memory profile and compare it