class_docstr: Collect and plot the memory stats at operator level.
  
  Includes ``memories_allocated``, ``memories_active`` and ``memories_reserved``.
  It also prints a summary for the top 20 operators that generate the most memories.
  
  Example usage:
  
      >>> # xdoctest: +SKIP(failing)
      >>> net.cuda()
      >>> input = input.cuda()
  
      >>> mem_tracker = MemoryTracker()
      >>> mem_tracker.start_monitor(net)
  
      >>> net.zero_grad(True)
      >>> loss = net(input)
      >>> if isinstance(loss, dict):
      >>>    loss = loss['out']
      >>> loss.sum().backward()
      >>> net.zero_grad(set_to_none=True)
  
      >>> mem_tracker.stop()
      >>> mem_tracker.summary()
      >>> mem_tracker.show_traces()
class_docstr_tok: ["Collect","and","plot","the","memory","stats","at","operator","level",".","Includes","``","memories_allocated","``",",","``","memories_active","``","and","``","memories_reserved","``",".","It","also","prints","a","summary","for","the","top","20","operators","that","generate","the","most","memories",".","Example","usage",":",">",">",">","#","xdoctest",":","+SKIP","(","failing",")",">",">",">","net.cuda","(",")",">",">",">","input","=","input.cuda","(",")",">",">",">","mem_tracker","=","MemoryTracker","(",")",">",">",">","mem_tracker.start_monitor","(","net",")",">",">",">","net.zero_grad","(","True",")",">",">",">","loss","=","net","(","input",")",">",">",">","if","isinstance","(","loss",",","dict",")",":",">",">",">","loss","=","loss","[","'out","'","]",">",">",">","loss.sum","(",")",".backward","(",")",">",">",">","net.zero_grad","(","set_to_none=True",")",">",">",">","mem_tracker.stop","(",")",">",">",">","mem_tracker.summary","(",")",">",">",">","mem_tracker.show_traces","(",")"]
docstr: Register module hooks and entering ``MemoryProfileDispatchMode``.
  
  This enables operator level memory stats can be tracked during module runtime.
docstr_tok: ["Register","module","hooks","and","entering","``","MemoryProfileDispatchMode","``",".","This","enables","operator","level","memory","stats","can","be","tracked","during","module","runtime","."]
code: def start_monitor(self, root_module: nn.Module) -> None:
      """
      Register module hooks and entering ``MemoryProfileDispatchMode``.
  
      This enables operator level memory stats can be tracked during module runtime.
      """
      self._clear_state()
      root_module.__setattr__("_memory_tracker_is_root", True)
      for name, m in root_module.named_modules():
          if m is not root_module:
              m.__setattr__("_memory_tracker_is_root", False)
          # fused_proxy_group does not support hooks
          if ".fused_proxy_grouped_embedding_bag" in name:
              continue
          # hook ordering with other hooks added by users is not managed, so
          # the memory stats tracked here may not completely accurate.
          h1 = m.register_forward_pre_hook(
              self._create_pre_forward_hook(name)
          )
          h2 = m.register_forward_hook(
              self._create_post_forward_hook(name)
          )
          # it does not work well with jagged tensor somehow, the root cause is not
          # clear and remove it for now as it does not really capture important info.
          # h3 = m.register_backward_hook(self._create_backward_hook(name))
          self._hooks.extend([h1, h2])
      torch.cuda.empty_cache()
      assert getattr(self, "profile_mode", None) is None
      self.profile_mode = MemoryProfileDispatchMode(self)
      self.profile_mode.__enter__()
  
imports: from collections import defaultdict
  from itertools import chain
  import pickle
  from typing import Any, Callable, Dict, List, no_type_check, Sequence
  import torch
  import torch.nn
  from torch.utils.hooks import RemovableHandle
  from torch.utils._python_dispatch import TorchDispatchMode
label_desc: Use image node_id 2 for calling the MemoryTracker obj's underlying member method code with example usage: obj.start_monitor(root_module) without return types