neural_compressor.experimental.export.torch2onnx

Helper functions to export model from PyTorch/TensorFlow to ONNX.

Module Contents

Functions

torch_to_fp32_onnx(pt_model, save_path, example_inputs)

Export FP32 PyTorch model into FP32 ONNX model.

torch_to_int8_onnx(pt_model, save_path, ...[, ...])

Export INT8 PyTorch model into INT8 ONNX model.
neural_compressor.experimental.export.torch2onnx.torch_to_fp32_onnx(pt_model, save_path, example_inputs, opset_version=14, dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}, input_names=None, output_names=None, do_constant_folding=True, verbose=True)[source]

Export FP32 PyTorch model into FP32 ONNX model.

Parameters:

  • pt_model (torch.nn.module) – PyTorch model.

  • save_path (str) – save path of ONNX model.

  • example_inputs (dict|list|tuple|torch.Tensor) – used to trace torch model.

  • opset_version (int, optional) – opset version. Defaults to 14.

  • dynamic_axes (dict, optional) – dynamic axes. Defaults to {“input”: {0: “batch_size”}, “output”: {0: “batch_size”}}.

  • input_names (dict, optional) – input names. Defaults to None.

  • output_names (dict, optional) – output names. Defaults to None.

  • do_constant_folding (bool, optional) – do constant folding or not. Defaults to True.

  • verbose (bool, optional) – dump verbose or not. Defaults to True.

neural_compressor.experimental.export.torch2onnx.torch_to_int8_onnx(pt_model, save_path, example_inputs, q_config, opset_version=14, dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}, input_names=None, output_names=None, quant_format: str = 'QDQ', verbose=True)[source]

Export INT8 PyTorch model into INT8 ONNX model.

Parameters:

  • pt_model (torch.nn.module) – PyTorch model.

  • save_path (str) – save path of ONNX model.

  • example_inputs (dict|list|tuple|torch.Tensor) – used to trace torch model.

  • q_config (dict) – containing quantization configuration.

  • opset_version (int, optional) – opset version. Defaults to 14.

  • dynamic_axes (dict, optional) – dynamic axes. Defaults to {“input”: {0: “batch_size”}, “output”: {0: “batch_size”}}.

  • input_names (dict, optional) – input names. Defaults to None.

  • output_names (dict, optional) – output names. Defaults to None.

  • quant_format (str, optional) – _quantization format of ONNX model. Defaults to ‘QDQ’.

  • verbose (bool, optional) – dump verbose or not. Defaults to True.