parameters: ()
docstring: main function
docstring_summary: main function
docstring_tokens: main function
function: def main():
      """
      main function
      """
      print("---- Tengine NanoDet-m Optimize Tool ----\n")
  
      args = parse_args()
      if args == None or args.input == None:
          usage_info()
          return None
  
      print(" Input model path    : %s" % (args.input))
      print("Output model path    : %s" % (args.output))
  
      # load original onnx model, graph, nodes
      print("[Opt Tools Info]: Step 0, load original onnx model from %s." % (args.input))
      onnx_model = onnx.load(args.input)
  
      if args.softmax:
          constant_shape_list = args.const.split(',')
  
          # update input constant nodes
          print("[Opt Tools Info]: Step 1, update the output shape in all dis_pred branches.")
          optimize_node_shape(onnx_model.graph.node, constant_shape_list)
  
          # add additional softmax nodes
          print("[Opt Tools Info]: Step 2, add Softmax node in the end of all dis_pred branche.")
          optimize_add_softmax(onnx_model.graph.node)
  
          # update output tensor name and shape
          print("[Opt Tools Info]: Step 3, update output tensor name and shape.")
          optimize_output_tensor(onnx_model.graph.output)
  
      # do check and simplify the onnx model
      print("[Opt Tools Info]: Step 4, check and simplify the new onnx model.")
      onnx_model, check = simplify(onnx_model)
  
      # save the new optimize onnx model
      print("[Opt Tools Info]: Step 5, save the new onnx model to %s" % (args.output))
      onnx.save(onnx_model, args.output)
  
      print("\n---- Tengine NanoDet-m Optimize onnx create success, best wish for your inference has a high accuracy ...\\(^0^)/ ----")
function_tokens: def main ( ) : print ( "---- Tengine NanoDet-m Optimize Tool ----\n" ) args = parse_args ( ) if args == None or args . input == None : usage_info ( ) return None print ( " Input model path    : %s" % ( args . input ) ) print ( "Output model path    : %s" % ( args . output ) ) # load original onnx model, graph, nodes print ( "[Opt Tools Info]: Step 0, load original onnx model from %s." % ( args . input ) ) onnx_model = onnx . load ( args . input ) if args . softmax : constant_shape_list = args . const . split ( ',' ) # update input constant nodes print ( "[Opt Tools Info]: Step 1, update the output shape in all dis_pred branches." ) optimize_node_shape ( onnx_model . graph . node , constant_shape_list ) # add additional softmax nodes print ( "[Opt Tools Info]: Step 2, add Softmax node in the end of all dis_pred branche." ) optimize_add_softmax ( onnx_model . graph . node ) # update output tensor name and shape print ( "[Opt Tools Info]: Step 3, update output tensor name and shape." ) optimize_output_tensor ( onnx_model . graph . output ) # do check and simplify the onnx model print ( "[Opt Tools Info]: Step 4, check and simplify the new onnx model." ) onnx_model , check = simplify ( onnx_model ) # save the new optimize onnx model print ( "[Opt Tools Info]: Step 5, save the new onnx model to %s" % ( args . output ) ) onnx . save ( onnx_model , args . output ) print ( "\n---- Tengine NanoDet-m Optimize onnx create success, best wish for your inference has a high accuracy ...\\(^0^)/ ----" )