docstr: ModelProfiler is used to share state to hooks.
  
  Parameters
  ----------
  custom_ops: dict
      a mapping of (module -> torch.nn.Module : custom operation)
      the custom operation is a callback funtion to calculate
      the module flops, parameters and the weight shape, it will overwrite the default operation.
      for reference, please see ``self.ops``.
  mode:
      the mode of how to collect information. If the mode is set to `default`,
      only the information of convolution, linear and rnn modules will be collected.
      If the mode is set to `full`, other operations will also be collected.
docstr_tok: ["ModelProfiler","is","used","to","share","state","to","hooks",".","Parameters","--","--","--","--","--","custom_ops",":","dict","a","mapping","of","(","module","-",">","torch.nn.Module",":","custom","operation",")","the","custom","operation","is","a","callback","funtion","to","calculate","the","module","flops",",","parameters","and","the","weight","shape",",","it","will","overwrite","the","default","operation",".","for","reference",",","please","see","``","self.ops","``",".","mode",":","the","mode","of","how","to","collect","information",".","If","the","mode","is","set","to","`","default","`",",","only","the","information","of","convolution",",","linear","and","rnn","modules","will","be","collected",".","If","the","mode","is","set","to","`","full","`",",","other","operations","will","also","be","collected","."]
code: def __init__(self, custom_ops=None, mode="default"):
      """
      ModelProfiler is used to share state to hooks.
  
      Parameters
      ----------
      custom_ops: dict
          a mapping of (module -> torch.nn.Module : custom operation)
          the custom operation is a callback funtion to calculate
          the module flops, parameters and the weight shape, it will overwrite the default operation.
          for reference, please see ``self.ops``.
      mode:
          the mode of how to collect information. If the mode is set to `default`,
          only the information of convolution, linear and rnn modules will be collected.
          If the mode is set to `full`, other operations will also be collected.
      """
      self.ops = {
          nn.Conv1d: self._count_convNd,
          nn.Conv2d: self._count_convNd,
          nn.Conv3d: self._count_convNd,
          nn.ConvTranspose1d: self._count_convNd,
          nn.ConvTranspose2d: self._count_convNd,
          nn.ConvTranspose3d: self._count_convNd,
          nn.Linear: self._count_linear,
          nn.RNNCell: self._count_rnn_cell,
          nn.GRUCell: self._count_gru_cell,
          nn.LSTMCell: self._count_lstm_cell,
          nn.RNN: self._count_rnn,
          nn.GRU: self._count_gru,
          nn.LSTM: self._count_lstm,
      }
      self._count_bias = False
      if mode == "full":
          self.ops.update(
              {
                  nn.BatchNorm1d: self._count_bn,
                  nn.BatchNorm2d: self._count_bn,
                  nn.BatchNorm3d: self._count_bn,
                  nn.LeakyReLU: self._count_relu,
                  nn.AvgPool1d: self._count_avgpool,
                  nn.AvgPool2d: self._count_avgpool,
                  nn.AvgPool3d: self._count_avgpool,
                  nn.AdaptiveAvgPool1d: self._count_adap_avgpool,
                  nn.AdaptiveAvgPool2d: self._count_adap_avgpool,
                  nn.AdaptiveAvgPool3d: self._count_adap_avgpool,
                  nn.Upsample: self._count_upsample,
                  nn.UpsamplingBilinear2d: self._count_upsample,
                  nn.UpsamplingNearest2d: self._count_upsample,
              }
          )
          self._count_bias = True
  
      if custom_ops is not None:
          self.ops.update(custom_ops)
  
      self.mode = mode
      self.results = []
  
imports: import functools
  from collections import Counter
  from prettytable import PrettyTable
  import torch
  import torch.nn
  from torch.nn.utils.rnn import PackedSequence
  from nni.compression.base.wrapper import ModuleWrapper
label_desc: Use image node_id 1 to create a new ModelProfiler object  with example: obj = ModelProfiler(custom_ops, mode)