code: def __init__(self, model, input_shape=(2, 3, 224, 224), device="cpu"):
      self.block_num = len(model.blocks)
      self.choice_num = len(model.blocks[0])
      self.flops_dict = {}
      self.params_dict = {}
  
      if device == "cpu":
          model = model.cpu()
      else:
          model = model.cuda()
  
      self.params_fixed = 0
      self.flops_fixed = 0
  
      input = torch.randn(input_shape)
  
      flops, params = get_model_complexity_info(
          model.conv_stem,
          (3, 224, 224),
          as_strings=False,
          print_per_layer_stat=False,
      )
      self.params_fixed += params / 1e6
      self.flops_fixed += flops / 1e6
  
      input = model.conv_stem(input)
  
      for block_id, block in enumerate(model.blocks):
          self.flops_dict[block_id] = {}
          self.params_dict[block_id] = {}
          for module_id, module in enumerate(block):
              flops, params = get_model_complexity_info(
                  module,
                  tuple(input.shape[1:]),
                  as_strings=False,
                  print_per_layer_stat=False,
              )
              # Flops(M)
              self.flops_dict[block_id][module_id] = flops / 1e6
              # Params(M)
              self.params_dict[block_id][module_id] = params / 1e6
  
          input = module(input)
  
      # conv_last
      flops, params = get_model_complexity_info(
          model.global_pool,
          tuple(input.shape[1:]),
          as_strings=False,
          print_per_layer_stat=False,
      )
      self.params_fixed += params / 1e6
      self.flops_fixed += flops / 1e6
  
      input = model.global_pool(input)
  
      # globalpool
      flops, params = get_model_complexity_info(
          model.conv_head,
          tuple(input.shape[1:]),
          as_strings=False,
          print_per_layer_stat=False,
      )
      self.params_fixed += params / 1e6
      self.flops_fixed += flops / 1e6
  
imports: import torch
  from ptflops import get_model_complexity_info
label_desc: Use image node_id 1 to create a new FlopsEst object from inherited base classes: object with example: obj = FlopsEst(model, input_shape, device)