code: def torch_to_onnx(
      model, config, input_shape, model_path, input_names, output_names
  ):
      """
      Convert torch model to onnx model and get layer bits config of onnx model.
  
      Parameters
      ----------
      model : pytorch model
          The model to speedup by quantization
      config : dict
          Config recording bits number and name of layers
      input_shape : tuple
          The input shape of model, shall pass it to torch.onnx.export
      model_path : str
          The path user want to store onnx model which is converted from pytorch model
      input_names : list
          Input name of onnx model providing for torch.onnx.export to generate onnx model
      output_name : list
          Output name of onnx model providing for torch.onnx.export to generate onnx model
  
      Returns
      -------
      onnx model
          Onnx model which is converted from pytorch model
      dict
          The configuration of onnx model layers and calibration parameters
      """
      # Support Gemm, Conv, Relu, Clip(Relu6) and MaxPool
      support_op = [
          torch.nn.Conv2d,
          torch.nn.Linear,
          torch.nn.ReLU,
          torch.nn.ReLU6,
          torch.nn.MaxPool2d,
      ]
      # Transfer bits number to onnx layer by using wrapper
      index2name = {}
      name2index = {}
      if config is not None:
          for i, name in enumerate(config.keys()):
              # i*2 to avoid the graph optimization of onnx on constant folding / elimination
              # for example, multiply 1 will be removed from the graph
              index2name[i * 2] = name
              name2index[name] = i * 2
      for name, module in model.named_modules():
          if config is not None and name in config:
              assert type(module) in support_op
              wrapper_module = LayernameModuleWrapper(
                  module, name2index[name]
              )
              set_nested_attr(model, name, wrapper_module)
          elif type(module) in support_op:
              wrapper_module = LayernameModuleWrapper(module, -1)
              set_nested_attr(model, name, wrapper_module)
      # Convert torch model to onnx model and save it in model_path
      device = torch.device("cpu")
      dummy_input = torch.randn(input_shape)
      dummy_input = dummy_input.to(device)
      model.to(device)
      torch.onnx.export(
          model,
          dummy_input,
          model_path,
          verbose=False,
          input_names=input_names,
          output_names=output_names,
          export_params=True,
      )
      # Load onnx model
      model_onnx = onnx.load(model_path)
      model_onnx, onnx_config = unwrapper(
          model_onnx, index2name, config
      )
      onnx.save(model_onnx, model_path)
  
      onnx.checker.check_model(model_onnx)
      return model_onnx, onnx_config
  
file_path: nni/nni/compression/quantization_speedup/frontend_to_onnx.py
imports: import torch
  import onnx
  import onnx.numpy_helper
  from ..utils import set_nested_attr
label_desc: Use image node_id 2 for calling a global function with example usage: torch_to_onnx(model, config, input_shape, model_path, input_names, output_names) and returns: model_onnx,  onnx_config