Skip to content

ONNX

model_navigator.api.config.OnnxConfig dataclass

Bases: CustomConfigForFormat

ONNX custom config used for ONNX export and conversion.

Parameters:

Name Type Description Default
opset Optional[int]

ONNX opset used for conversion.

DEFAULT_ONNX_OPSET
dynamic_axes Optional[Dict[str, Union[Dict[int, str], List[int]]]]

Dynamic axes for ONNX conversion.

None
onnx_extended_conversion bool

Enables additional conversions from TorchScript to ONNX.

False

format: Format property

Format represented by CustomConfig.

Returns:

Type Description
Format

OnnxConfig format

name() classmethod

Name of the config.

Source code in model_navigator/api/config.py
@classmethod
def name(cls) -> str:
    """Name of the config."""
    return "Onnx"

model_navigator.api.onnx

ONNX optimize API.

optimize(model, dataloader, sample_count=DEFAULT_SAMPLE_COUNT, batching=True, target_formats=None, runners=None, profiler_config=None, workspace=None, verbose=False, debug=False, verify_func=None, custom_configs=None)

Function exports ONNX model to all supported formats.

Parameters:

Name Type Description Default
model Union[Path, str]

ONNX model path or string

required
dataloader SizedDataLoader

Sized iterable with data that will be feed to the model

required
sample_count int

Limits how many samples will be used from dataloader

DEFAULT_SAMPLE_COUNT
batching Optional[bool]

Enable or disable batching on first (index 0) dimension of the model

True
target_formats Optional[Union[Union[str, Format], Tuple[Union[str, Format], ...]]]

Target model formats for optimize process

None
runners Optional[Union[Union[str, Type[NavigatorRunner]], Tuple[Union[str, Type[NavigatorRunner]], ...]]]

Use only runners provided as parameter

None
profiler_config Optional[ProfilerConfig]

Profiling config

None
workspace Optional[Path]

Workspace where packages will be extracted

None
verbose bool

Enable verbose logging

False
debug bool

Enable debug logging from commands

False
verify_func Optional[VerifyFunction]

Function for additional model verification

None
custom_configs Optional[Sequence[CustomConfig]]

Sequence of CustomConfigs used to control produced artifacts

None

Returns:

Type Description
Package

Package descriptor representing created package.

Source code in model_navigator/api/onnx.py
def optimize(
    model: Union[Path, str],
    dataloader: SizedDataLoader,
    sample_count: int = DEFAULT_SAMPLE_COUNT,
    batching: Optional[bool] = True,
    target_formats: Optional[Union[Union[str, Format], Tuple[Union[str, Format], ...]]] = None,
    runners: Optional[Union[Union[str, Type[NavigatorRunner]], Tuple[Union[str, Type[NavigatorRunner]], ...]]] = None,
    profiler_config: Optional[ProfilerConfig] = None,
    workspace: Optional[Path] = None,
    verbose: bool = False,
    debug: bool = False,
    verify_func: Optional[VerifyFunction] = None,
    custom_configs: Optional[Sequence[CustomConfig]] = None,
) -> Package:
    """Function exports ONNX model to all supported formats.

    Args:
        model: ONNX model path or string
        dataloader: Sized iterable with data that will be feed to the model
        sample_count: Limits how many samples will be used from dataloader
        batching: Enable or disable batching on first (index 0) dimension of the model
        target_formats: Target model formats for optimize process
        runners: Use only runners provided as parameter
        profiler_config: Profiling config
        workspace: Workspace where packages will be extracted
        verbose: Enable verbose logging
        debug: Enable debug logging from commands
        verify_func: Function for additional model verification
        custom_configs: Sequence of CustomConfigs used to control produced artifacts

    Returns:
        Package descriptor representing created package.
    """
    if isinstance(model, str):
        model = Path(model)
    if workspace is None:
        workspace = get_default_workspace()
    if target_formats is None:
        target_formats = DEFAULT_ONNX_TARGET_FORMATS

    if runners is None:
        runners = default_runners()

    if profiler_config is None:
        profiler_config = ProfilerConfig()

    target_formats_enums = enums.parse(target_formats, Format)
    runner_names = enums.parse(runners, lambda runner: runner if isinstance(runner, str) else runner.name())

    if Format.ONNX not in target_formats_enums:
        target_formats_enums = (Format.ONNX,) + target_formats_enums

    config = CommonConfig(
        Framework.ONNX,
        model=model,
        dataloader=dataloader,
        workspace=workspace,
        target_formats=target_formats_enums,
        sample_count=sample_count,
        batch_dim=0 if batching else None,
        runner_names=runner_names,
        profiler_config=profiler_config,
        verbose=verbose,
        debug=debug,
        verify_func=verify_func,
        custom_configs=map_custom_configs(custom_configs=custom_configs),
    )

    models_config = ModelConfigBuilder.generate_model_config(
        framework=Framework.ONNX,
        target_formats=target_formats_enums,
        custom_configs=custom_configs,
    )

    builders = [
        preprocessing_builder,
        onnx_export_builder,
        onnx_conversion_builder,
        correctness_builder,
    ]
    if profiler_config.run_profiling:
        builders.append(profiling_builder)
    builders.append(verify_builder)

    package = PipelineManager.run(
        pipeline_builders=builders,
        config=config,
        models_config=models_config,
    )

    return package