pytorch/docs/source/module_tracker.rst

torch.utils.module_tracker
===================================
.. automodule:: torch.utils.module_tracker

This utility can be used to track the current position inside an :class:`torch.nn.Module` hierarchy.
It can be used within other tracking tools to be able to easily associate measured quantities to user-friendly names. This is used in particular in the FlopCounterMode today.

.. autoclass:: torch.utils.module_tracker.ModuleTracker
Add module tracker (#125352) This does a few things that were originally a few PRs but I am on a new machine and don't have ghstack. If it is too problematic to review, I can re-split, just let me know. This does: - Cleanup context manager use in test_flop_counter - Remove need for mod argument in FlopCounterMode, warning about it - Re-implement a Module tracker from scratch using global forward Module use and multi_grad_hook (we cannot use global backward Module hook because they don't look for nested Tensor and they're custom Function based instead of multi_grad_hook). - Update FlopCouterMode to use the new ModuleTracker. All the existing test suite passes as-is (only changes there are new tests and refactoring mentioned above) Pull Request resolved: https://github.com/pytorch/pytorch/pull/125352 Approved by: https://github.com/mikaylagawarecki 2024-05-04 18:33:35 +00:00			`torch.utils.module_tracker`
			`===================================`
			`.. automodule:: torch.utils.module_tracker`

			This utility can be used to track the current position inside an :class:`torch.nn.Module` hierarchy.
			`It can be used within other tracking tools to be able to easily associate measured quantities to user-friendly names. This is used in particular in the FlopCounterMode today.`

			`.. autoclass:: torch.utils.module_tracker.ModuleTracker`