class_docstr: Observe a module to get the graph of the memory consumption during
  the forward and backward, layer by layer, with:
  - a breakdown of the memory used (activations memory estimation)
  - additional details such as amount of data exchanged with all gather
  
  Requires the model to be on a CUDA device to track its memory
  
  Example usage (no FSDP):
  
      ```
      # create your model
      model = models.resnet50().cuda()
  
      # monitor the model
      tracker = LayerwiseMemoryTracker()
      tracker.monitor(model)
  
      # Do a forward/backward
      criterion(model(input), target).backward()
  
      # show the plots
      tracker.show_plots()
  
      # get the detailed traces
      tracker.memory_traces
  
      # print a summary
      print(tracker.summary)
      ```
  
  Advanced usage (for FSDP):
  
      ```
      # wrap the group used for FSDP
      group = ProcessGroupTracker(group)
  
      # use this group when creating FSDP blocks
      model = FullyShardedDataParallel(model, process_group=group),
  
      # monitor the model as before
      tracker = LayerwiseMemoryTracker()
      tracker.monitor(model)
  
      # the detailed traces will now contain information
      # about the amount of all gathered data
      tracker.memory_traces
      ```
class_docstr_tok: ["Observe","a","module","to","get","the","graph","of","the","memory","consumption","during","the","forward","and","backward",",","layer","by","layer",",","with",":","-","a","breakdown","of","the","memory","used","(","activations","memory","estimation",")","-","additional","details","such","as","amount","of","data","exchanged","with","all","gather","Requires","the","model","to","be","on","a","CUDA","device","to","track","its","memory","Example","usage","(","no","FSDP",")",":","``","`","#","create","your","model","model","=","models.resnet50","(",")",".cuda","(",")","#","monitor","the","model","tracker","=","LayerwiseMemoryTracker","(",")","tracker.monitor","(","model",")","#","Do","a","forward\/backward","criterion","(","model","(","input",")",",","target",")",".backward","(",")","#","show","the","plots","tracker.show_plots","(",")","#","get","the","detailed","traces","tracker.memory_traces","#","print","a","summary","print","(","tracker.summary",")","``","`","Advanced","usage","(","for","FSDP",")",":","``","`","#","wrap","the","group","used","for","FSDP","group","=","ProcessGroupTracker","(","group",")","#","use","this","group","when","creating","FSDP","blocks","model","=","FullyShardedDataParallel","(","model",",","process_group=group",")",",","#","monitor","the","model","as","before","tracker","=","LayerwiseMemoryTracker","(",")","tracker.monitor","(","model",")","#","the","detailed","traces","will","now","contain","information","#","about","the","amount","of","all","gathered","data","tracker.memory_traces","``","`"]
docstr: Get the part of the traces which corresponds to the forward pass
docstr_tok: ["Get","the","part","of","the","traces","which","corresponds","to","the","forward","pass"]
code: def forward_traces(self) -> List[LayerMemoryTrace]:
      """
      Get the part of the traces which corresponds to the forward pass
      """
      return [t for t in self.memory_traces if t.is_forward]
  
imports: from contextlib import contextmanager
  from dataclasses import dataclass
  from enum import auto, Enum
  from functools import lru_cache
  from typing import Dict, List, NamedTuple, Tuple, Union
  import numpy
  import torch
  import torch.nn
  from fairscale.nn import FullyShardedDataParallel
  from vissl.utils.visualize import matplotlib_figure_to_image
label_desc: Use image node_id 5 for calling the LayerwiseMemoryTracker obj's underlying member method code with example usage: obj.forward_traces() and returns: unknown