code: def count_flops(sym, **data_shapes):
      all_layers = sym.get_internals()
      # print(all_layers)
      arg_shapes, out_shapes, aux_shapes = all_layers.infer_shape(
          **data_shapes
      )
      out_shape_dict = dict(zip(all_layers.list_outputs(), out_shapes))
  
      nodes = json.loads(sym.tojson())["nodes"]
      nodeid_shape = {}
      for nodeid, node in enumerate(nodes):
          name = node["name"]
          layer_name = name + "_output"
          if layer_name in out_shape_dict:
              nodeid_shape[nodeid] = out_shape_dict[layer_name]
      # print(nodeid_shape)
      FLOPs = 0
      for nodeid, node in enumerate(nodes):
          flops = 0
          if node["op"] == "Convolution":
              output_shape = nodeid_shape[nodeid]
              name = node["name"]
              attr = node["attrs"]
              input_nodeid = node["inputs"][0][0]
              input_shape = nodeid_shape[input_nodeid]
              flops = count_conv_flops(input_shape, output_shape, attr)
          elif node["op"] == "FullyConnected":
              attr = node["attrs"]
              output_shape = nodeid_shape[nodeid]
              input_nodeid = node["inputs"][0][0]
              input_shape = nodeid_shape[input_nodeid]
              output_filter = output_shape[1]
              input_filter = (
                  input_shape[1] * input_shape[2] * input_shape[3]
              )
              # assert len(input_shape)==4 and input_shape[2]==1 and input_shape[3]==1
              flops = count_fc_flops(input_filter, output_filter, attr)
          # print(node, flops)
          FLOPs += flops
  
      return FLOPs
  
file_path: insightface/recognition/arcface_mxnet/common/flops_counter.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import sys
  import os
  import json
  import argparse
  import numpy
  import mxnet
label_desc: Use image node_id 4 for calling a global function with example usage: count_flops(sym) and returns: FLOPs