PyTorch

model_navigator.api.config.JitType

Bases: Enum

TorchScript export paramter.

model_navigator.api.config.TorchConfig dataclass

Bases: CustomConfigForFormat

Torch custom config used for TorchScript models export.

Parameters:

Name	Type	Description	Default
jit_type	Union[Union[str, JitType], Tuple[Union[str, JitType], ...]]	Type of TorchScript export.	(JitType.SCRIPT, JitType.TRACE)

format: Format property

Format represented by CustomConfig.

Returns:

Type	Description
Format	TorchConfig format

__post_init__()

Parse dataclass enums.

Source code in model_navigator/api/config.py

def __post_init__(self) -> None:
    """Parse dataclass enums."""
    jit_type = (self.jit_type,) if not isinstance(self.jit_type, (list, tuple)) else self.jit_type
    self.jit_type = tuple(JitType(j) for j in jit_type)

defaults()

Update parameters to defaults.

Source code in model_navigator/api/config.py

def defaults(self) -> None:
    """Update parameters to defaults."""
    self.jit_type = (JitType.SCRIPT, JitType.TRACE)

name() classmethod

Name of the config.

Source code in model_navigator/api/config.py

@classmethod
def name(cls) -> str:
    """Name of the config."""
    return "Torch"

model_navigator.api.config.TorchTensorRTConfig dataclass

Bases: CustomConfigForFormat

Torch custom config used for TensorRT TorchScript conversion.

Parameters:

Name	Type	Description	Default
precision	Union[Union[str, TensorRTPrecision], Tuple[Union[str, TensorRTPrecision], ...]]	TensorRT precision.	DEFAULT_TENSORRT_PRECISION
max_workspace_size	Optional[int]	Max workspace size used by converter.	DEFAULT_MAX_WORKSPACE_SIZE
trt_profile	Optional[TensorRTProfile]	TensorRT profile.	None

format: Format property

Format represented by CustomConfig.

Returns:

Type	Description
Format	TorchTensorRTConfig format

__post_init__()

Parse dataclass enums.

Source code in model_navigator/api/config.py

def __post_init__(self) -> None:
    """Parse dataclass enums."""
    precision = (self.precision,) if not isinstance(self.precision, (list, tuple)) else self.precision
    self.precision = tuple(TensorRTPrecision(p) for p in precision)
    self.precision_mode = TensorRTPrecisionMode(self.precision_mode)

defaults()

Update parameters to defaults.

Source code in model_navigator/api/config.py

def defaults(self) -> None:
    """Update parameters to defaults."""
    self.precision = tuple(TensorRTPrecision(p) for p in DEFAULT_TENSORRT_PRECISION)
    self.precision_mode = DEFAULT_TENSORRT_PRECISION_MODE
    self.trt_profile = None
    self.max_workspace_size = DEFAULT_MAX_WORKSPACE_SIZE

from_dict(config_dict) classmethod

Instantiate TorchTensorRTConfig from adictionary.

Source code in model_navigator/api/config.py

@classmethod
def from_dict(cls, config_dict: Dict[str, Any]) -> "TorchTensorRTConfig":
    """Instantiate TorchTensorRTConfig from  adictionary."""
    if "trt_profile" in config_dict and not isinstance(config_dict["trt_profile"], TensorRTProfile):
        config_dict["trt_profile"] = TensorRTProfile.from_dict(config_dict["trt_profile"])
    return cls(**config_dict)

name() classmethod

Name of the config.

Source code in model_navigator/api/config.py

@classmethod
def name(cls) -> str:
    """Name of the config."""
    return "TorchTensorRT"

model_navigator.api.torch

Torch optimize API.

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

Function exports PyTorch model to all supported formats.

Parameters:

Name	Type	Description	Default
model	torch.nn.Module	PyTorch model object	required
dataloader	SizedDataLoader	Sized iterable with data that will be feed to the model	required
sample_count	Optional[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
input_names	Optional[Tuple[str, ...]]	Model input names	None
output_names	Optional[Tuple[str, ...]]	Model output names	None
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	Optional[bool]	Enable verbose logging	False
debug	Optional[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/torch.py

def optimize(
    model: torch.nn.Module,
    dataloader: SizedDataLoader,
    sample_count: Optional[int] = DEFAULT_SAMPLE_COUNT,
    batching: Optional[bool] = True,
    input_names: Optional[Tuple[str, ...]] = None,
    output_names: Optional[Tuple[str, ...]] = None,
    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: Optional[bool] = False,
    debug: Optional[bool] = False,
    verify_func: Optional[VerifyFunction] = None,
    custom_configs: Optional[Sequence[CustomConfig]] = None,
) -> Package:
    """Function exports PyTorch model to all supported formats.

    Args:
        model: PyTorch model object
        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
        input_names: Model input names
        output_names: Model output names
        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 workspace is None:
        workspace = get_default_workspace()
    if target_formats is None:
        target_formats = DEFAULT_TORCH_TARGET_FORMATS
        if batching:
            target_formats, custom_configs = update_allowed_batching_parameters(
                target_formats=target_formats,
                custom_configs=custom_configs,
            )
        LOGGER.info(f"Using default target formats: {[tf.name for tf in target_formats]}")

    target_device = "cuda" if torch.cuda.is_available() else "cpu"

    sample = next(iter(dataloader))
    if isinstance(sample, Mapping):
        forward_kw_names = tuple(sample.keys())
    else:
        forward_kw_names = None

    if runners is None:
        runners = default_runners()

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

    if profiler_config is None:
        profiler_config = ProfilerConfig()

    if Format.TORCH not in target_formats:
        target_formats = (Format.TORCH,) + target_formats

    config = CommonConfig(
        framework=Framework.TORCH,
        model=model,
        dataloader=dataloader,
        target_formats=target_formats,
        workspace=workspace,
        sample_count=sample_count,
        _input_names=input_names,
        _output_names=output_names,
        target_device=target_device,
        forward_kw_names=forward_kw_names,
        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.TORCH,
        target_formats=target_formats,
        custom_configs=custom_configs,
    )

    builders = [
        preprocessing_builder,
        torch_export_builder,
        torch_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