torchtnt.utils.module_summary.get_module_summary

torchtnt.utils.module_summary.get_module_summary(module: Module, module_args: Optional[Tuple[Any, ...]] = None, module_kwargs: Optional[MutableMapping[str, Any]] = None) → ModuleSummary

Generate a ModuleSummary object, then assign its values and generate submodule tree.

Parameters:

• module – The module to be summarized.
• module_args – A tuple of arguments for the module to run and calculate FLOPs and activation sizes.
• module_kwargs –

  Any kwarg arguments to be passed into the module’s forward function.

  Note

  To calculate FLOPs, you must use PyTorch 1.13 or greater.

  Note

  If module contains any lazy submodule, we will NOT calculate FLOPs.

  Note

  Currently only modules that output a single tensor are supported. TODO: to support more flexible output for module.