Skip to main content
Version: v2509

convert.converter.convert_onnx

convert_trt_with_onnx​

@register_conversion("onnx")
def convert_trt_with_onnx(model: nn.Module,
                          input_args: Sequence[Tuple[Tuple[Any], Dict[str,
                                                                      Any]]],
                          export_path: str,
                          int8: bool = False,
                          fp16: bool = False,
                          calibrator: Union[Calibrator, None] = None,
                          input_names: Union[List[str], None] = None,
                          dynamic_axes: Union[Dict[str,
                                                   Dict[int,
                                                        Union[str,
                                                              Tuple[str, int,
                                                                    int]]]],
                                              None] = None,
                          **kwargs) -> Dict[str, Any]

Converts a PyTorch model to ONNX format and then to TensorRT engine using ONNX GraphSurgeon.

Exports the model to ONNX, sanitizes it with onnx_graphsurgeon, and builds a TensorRT engine. Supports INT8/FP16 precision and dynamic axes configuration.

Arguments:

  • model nn.Module - The PyTorch model to convert.
  • input_args Sequence[Tuple[Tuple[Any], Dict[str, Any]]] - Input arguments for model tracing.
  • export_path str - Path to save the TensorRT engine file.
  • int8 bool, optional - Enable INT8 precision. Defaults to False.
  • fp16 bool, optional - Enable FP16 precision. Defaults to False.
  • calibrator Calibrator | None, optional - Custom calibrator for INT8. Defaults to None.
  • input_names List[str] | None, optional - Names of input tensors. Defaults to None. dynamic_axes (Dict[str, Dict[int, str | Tuple[str, int, int]]] | None, optional): Dictionary specifying dynamic axes for ONNX export. Defaults to None.
  • **kwargs - Additional keyword arguments passed to torch.onnx.export.

Returns:

Dict[str, Any]: Conversion metadata including:

  • 'rt_mode': Always 'onnx'
  • 'input_shapes': Input tensor shapes
  • 'int8': Whether INT8 was enabled
  • 'fp16': Whether FP16 was enabled
  • 'input_names': Input tensor names if provided
  • 'dynamic_axes': Dynamic axes configuration if provided

Raises:

  • AssertionError - If input_names is missing when dynamic_axes is provided.