Config

Classes, enums and types used to configure Model Navigator.

model_navigator.api.config

Definition of enums and classes representing configuration for Model Navigator.

CustomConfig

Bases: abc.ABC

Base class used for custom configs. Input for Model Navigator optimize method.

defaults

defaults()

Update parameters to defaults.

Source code in model_navigator/api/config.py

def defaults(self) -> None:
    """Update parameters to defaults."""
    return None

from_dict classmethod

from_dict(config_dict)

Instantiate CustomConfig from a dictionary.

Source code in model_navigator/api/config.py

@classmethod
def from_dict(cls, config_dict: Dict[str, Any]) -> "CustomConfig":
    """Instantiate CustomConfig from a dictionary."""
    return cls(**config_dict)

name abstractmethod classmethod

name()

Name of the CustomConfig.

Source code in model_navigator/api/config.py

@classmethod
@abc.abstractmethod
def name(cls) -> str:
    """Name of the CustomConfig."""
    raise NotImplementedError()

CustomConfigForFormat

Bases: DataObject, CustomConfig

Abstract base class used for custom configs representing particular format.

format property abstractmethod

format: Format

Format represented by CustomConfig.

DeviceKind

Bases: Enum

Supported types of devices.

Parameters:

• CPU (str) –

  Select CPU device.

• GPU (str) –

  Select GPU with CUDA support.

Format

Bases: Enum

All model formats supported by Model Navigator 'optimize' function.

Parameters:

• PYTHON (str) –

  Format indicating any model defined in Python.

• TORCH (str) –

  Format indicating PyTorch model.

• TENSORFLOW (str) –

  Format indicating TensorFlow model.

• JAX (str) –

  Format indicating JAX model.

• TORCHSCRIPT (str) –

  Format indicating TorchScript model.

• TF_SAVEDMODEL (str) –

  Format indicating TensorFlow SavedModel.

• TF_TRT (str) –

  Format indicating TensorFlow TensorRT model.

• TORCH_TRT (str) –

  Format indicating PyTorch TensorRT model.

• ONNX (str) –

  Format indicating ONNX model.

• TENSORRT (str) –

  Format indicating TensorRT model.

JitType

Bases: Enum

TorchScript export parameter.

Used for selecting the type of TorchScript export.

Parameters:

• TRACE (str) –

  Use tracing during export.

• SCRIPT (str) –

  Use scripting during export.

OnnxConfig dataclass

Bases: CustomConfigForFormat

ONNX custom config used for ONNX export and conversion.

Parameters:

• opset (Optional[int]) –

  ONNX opset used for conversion.

• dynamic_axes (Optional[Dict[str, Union[Dict[int, str], List[int]]]]) –

  Dynamic axes for ONNX conversion.

• onnx_extended_conversion (bool) –

  Enables additional conversions from TorchScript to ONNX.

format property

format: Format

Returns Format.ONNX.

Returns:

• Format –

  Format.ONNX

name classmethod

name()

Name of the config.

Source code in model_navigator/api/config.py

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

OptimizationProfile dataclass

Bases: DataObject

Optimization profile configuration.

For each batch size profiler will run measurements in windows of fixed number of queries. Batch sizes are profiled in the ascending order.

Profiler will run multiple trials and will stop when the measurements are stable (within stability_percentage from the mean) within three consecutive windows. If the measurements are not stable after max_trials trials, the profiler will stop with an error. Profiler will also stop profiling when the throughput does not increase at least by throughput_cutoff_threshold.

Parameters:

• max_batch_size (Optional[int]) –

  Maximal batch size used during conversion and profiling. None mean automatic search is enabled.

• batch_sizes –

  List of batch sizes to profile. None mean automatic search is enabled.

• window_size (Optional[int]) –

  Number of requests to measure in each window.

• stability_percentage (float) –

  Allowed percentage of variation from the mean in three consecutive windows.

• max_trials (int) –

  Maximum number of window trials.

• throughput_cutoff_threshold (float) –

  Minimum throughput increase to continue profiling.

• dataloader (Optional[SizedDataLoader]) –

  Optional dataloader for profiling. Use only 1 sample.

from_dict classmethod

from_dict(optimization_profile_dict)

Instantiate OptimizationProfile class from a dictionary.

Parameters:

• optimization_profile_dict (Mapping) –

  Data dictionary.

Returns:

• OptimizationProfile –

  OptimizationProfile

Source code in model_navigator/api/config.py

@classmethod
def from_dict(cls, optimization_profile_dict: Mapping) -> "OptimizationProfile":
    """Instantiate OptimizationProfile class from a dictionary.

    Args:
        optimization_profile_dict (Mapping): Data dictionary.

    Returns:
        OptimizationProfile
    """
    return cls(
        max_batch_size=optimization_profile_dict.get("max_batch_size"),
        batch_sizes=optimization_profile_dict.get("batch_sizes"),
        window_size=optimization_profile_dict.get("window_size"),
        stability_percentage=optimization_profile_dict.get("stability_percentage", 10.0),
        max_trials=optimization_profile_dict.get("max_trials", 10),
        throughput_cutoff_threshold=optimization_profile_dict.get("throughput_cutoff_threshold", -2),
    )

to_dict

to_dict(filter_fields=None, parse=False)

Serialize to a dictionary.

Append dataloader field to filtered fields during dump.

Parameters:

• filter_fields (Optional[List[str]]) –

  List of fields to filter out. Defaults to None.

• parse (bool) –

  If True recursively parse field values to jsonable representation. Defaults to False.

Returns:

• Dict( Dict ) –

  Data serialized to a dictionary.

Source code in model_navigator/api/config.py

def to_dict(self, filter_fields: Optional[List[str]] = None, parse: bool = False) -> Dict:
    """Serialize to a dictionary.

    Append `dataloader` field to filtered fields during dump.

    Args:
        filter_fields (Optional[List[str]], optional): List of fields to filter out.
            Defaults to None.
        parse (bool, optional): If True recursively parse field values to jsonable representation.
            Defaults to False.

    Returns:
        Dict: Data serialized to a dictionary.
    """
    if not filter_fields:
        filter_fields = []

    filter_fields += ["dataloader"]
    return super().to_dict(filter_fields=filter_fields, parse=parse)

ShapeTuple dataclass

Bases: DataObject

Represents a set of shapes for a single binding in a profile.

Each element of the tuple represents a shape for a single dimension of the binding.

Parameters:

• min (Tuple[int]) –

  The minimum shape that the profile will support.

• opt (Tuple[int]) –

  The shape for which TensorRT will optimize the engine.

• max (Tuple[int]) –

  The maximum shape that the profile will support.

__iter__

__iter__()

Iterate over shapes.

Source code in model_navigator/api/config.py

def __iter__(self):
    """Iterate over shapes."""
    yield from [self.min, self.opt, self.max]

__repr__

__repr__()

Representation.

Source code in model_navigator/api/config.py

def __repr__(self):
    """Representation."""
    return type(self).__name__ + self.__str__()

__str__

__str__()

String representation.

Source code in model_navigator/api/config.py

def __str__(self):
    """String representation."""
    return f"(min={self.min}, opt={self.opt}, max={self.max})"

SizedIterable

Bases: Protocol

Protocol representing sized iterable. Used by dataloader.

__iter__

__iter__()

Magic method iter.

Returns:

• Iterator –

  Iterator to next item.

Source code in model_navigator/api/config.py

def __iter__(self) -> Iterator:
    """Magic method __iter__.

    Returns:
        Iterator to next item.
    """
    ...

__len__

__len__()

Magic method len.

Returns:

• int –

  Length of size iterable.

Source code in model_navigator/api/config.py

def __len__(self) -> int:
    """Magic method __len__.

    Returns:
        Length of size iterable.
    """
    ...

TensorFlowConfig dataclass

Bases: CustomConfigForFormat

TensorFlow custom config used for SavedModel export.

Parameters:

• jit_compile (Tuple[Optional[bool], ...]) –

  Enable or Disable jit_compile flag for tf.function wrapper for Jax infer function.

• enable_xla (Tuple[Optional[bool], ...]) –

  Enable or Disable enable_xla flag for jax2tf converter.

format property

format: Format

Returns Format.TF_SAVEDMODEL.

Returns:

• Format –

  Format.TF_SAVEDMODEL

defaults

defaults()

Update parameters to defaults.

Source code in model_navigator/api/config.py

def defaults(self) -> None:
    """Update parameters to defaults."""
    self.jit_compile = (None,)
    self.enable_xla = (None,)

name classmethod

name()

Name of the config.

Source code in model_navigator/api/config.py

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

TensorFlowTensorRTConfig dataclass

Bases: CustomConfigForFormat

TensorFlow TensorRT custom config used for TensorRT SavedModel export.

Parameters:

• precision (Union[Union[str, TensorRTPrecision], Tuple[Union[str, TensorRTPrecision], ...]]) –

  TensorRT precision.

• max_workspace_size (Optional[int]) –

  Max workspace size used by converter.

• minimum_segment_size (int) –

  Min size of subgraph.

• trt_profile (Optional[TensorRTProfile]) –

  TensorRT profile.

format property

format: Format

Returns Format.TF_TRT.

Returns:

• Format –

  Format.TF_TRT

__post_init__

__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)

defaults

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.max_workspace_size = DEFAULT_MAX_WORKSPACE_SIZE
    self.minimum_segment_size = DEFAULT_MIN_SEGMENT_SIZE
    self.trt_profile = None

from_dict classmethod

from_dict(config_dict)

Instantiate TensorFlowTensorRTConfig from adictionary.

Source code in model_navigator/api/config.py

@classmethod
def from_dict(cls, config_dict: Dict[str, Any]) -> "TensorFlowTensorRTConfig":
    """Instantiate TensorFlowTensorRTConfig from  adictionary."""
    if config_dict.get("trt_profile") is not None 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()

Name of the config.

Source code in model_navigator/api/config.py

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

TensorRTCompatibilityLevel

Bases: Enum

Compatibility level for TensorRT.

Parameters:

• AMPERE_PLUS (str) –

  Support AMPERE plus architecture

TensorRTConfig dataclass

Bases: CustomConfigForFormat

TensorRT custom config used for TensorRT conversion.

Parameters:

• precision (Union[Union[str, TensorRTPrecision], Tuple[Union[str, TensorRTPrecision], ...]]) –

  TensorRT precision.

• max_workspace_size (Optional[int]) –

  Max workspace size used by converter.

• trt_profile (Optional[TensorRTProfile]) –

  TensorRT profile.

• optimization_level (Optional[int]) –

  Optimization level for TensorRT conversion. Allowed values are fom 0 to 5. Where default is 3 based on TensorRT API documentation.

format property

format: Format

Returns Format.TENSORRT.

Returns:

• Format –

  Format.TENSORRT

__post_init__

__post_init__()

Parse dataclass enums.

Source code in model_navigator/api/config.py

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

    if self.optimization_level is not None and (self.optimization_level < 0 or self.optimization_level > 5):
        raise ModelNavigatorConfigurationError(
            f"TensorRT `optimization_level` must be between 0 and 5. Provided value: {self.optimization_level}."
        )

defaults

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
    self.optimization_level = None
    self.compatibility_level = None

from_dict classmethod

from_dict(config_dict)

Instantiate TensorRTConfig from adictionary.

Source code in model_navigator/api/config.py

@classmethod
def from_dict(cls, config_dict: Dict[str, Any]) -> "TensorRTConfig":
    """Instantiate TensorRTConfig from  adictionary."""
    if config_dict.get("trt_profile") is not None 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()

Name of the config.

Source code in model_navigator/api/config.py

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

TensorRTPrecision

Bases: Enum

Precisions supported during TensorRT conversions.

Parameters:

• INT8 (str) –

  8-bit integer precision.

• FP16 (str) –

  16-bit floating point precision.

• FP32 (str) –

  32-bit floating point precision.

TensorRTPrecisionMode

Bases: Enum

Precision modes for TensorRT conversions.

Parameters:

• HIERARCHY (str) –

  Use TensorRT precision hierarchy starting from highest to lowest.

• SINGLE (str) –

  Use single precision.

• MIXED (str) –

  Use mixed precision.

TensorRTProfile

Bases: Dict[str, ShapeTuple]

Single optimization profile that can be used to build an engine.

More specifically, it is an Dict[str, ShapeTuple] which maps binding names to a set of min/opt/max shapes.

__getitem__

__getitem__(key)

Retrieves the shapes registered for a given input name.

Returns:

• ShapeTuple –

  A named tuple including ``min``, ``opt``, and ``max`` members for the shapes
  corresponding to the input.
  

Source code in model_navigator/api/config.py

def __getitem__(self, key):
    """Retrieves the shapes registered for a given input name.

    Returns:
        ShapeTuple:
                A named tuple including ``min``, ``opt``, and ``max`` members for the shapes
                corresponding to the input.
    """
    if key not in self:
        LOGGER.error(f"Binding: {key} does not have shapes set in this profile")
    return super().__getitem__(key)

__repr__

__repr__()

Representation.

Source code in model_navigator/api/config.py

def __repr__(self):
    """Representation."""
    ret = "TensorRTProfile()"
    for name, (min, opt, max) in self.items():
        ret += f".add('{name}', min={min}, opt={opt}, max={max})"
    return ret

__str__

__str__()

String representation.

Source code in model_navigator/api/config.py

def __str__(self):
    """String representation."""
    elems = []
    for name, (min, opt, max) in self.items():
        elems.append(f"{name} [min={min}, opt={opt}, max={max}]")

    sep = ",\n "
    return "{" + sep.join(elems) + "}"

add

add(name, min, opt, max)

A convenience function to add shapes for a single binding.

Parameters:

• name (str) –

  The name of the binding.

• min (Tuple[int]) –

  The minimum shape that the profile will support.

• opt (Tuple[int]) –

  The shape for which TensorRT will optimize the engine.

• max (Tuple[int]) –

  The maximum shape that the profile will support.

Returns:

• Profile –

  self, which allows this function to be easily chained to add multiple bindings, e.g., TensorRTProfile().add(...).add(...)

Source code in model_navigator/api/config.py

def add(self, name, min, opt, max):
    """A convenience function to add shapes for a single binding.

    Args:
        name (str): The name of the binding.
        min (Tuple[int]): The minimum shape that the profile will support.
        opt (Tuple[int]): The shape for which TensorRT will optimize the engine.
        max (Tuple[int]): The maximum shape that the profile will support.

    Returns:
        Profile:
            self, which allows this function to be easily chained to add multiple bindings,
            e.g., TensorRTProfile().add(...).add(...)
    """
    self[name] = ShapeTuple(min, opt, max)
    return self

from_dict classmethod

from_dict(profile_dict)

Create a TensorRTProfile from a dictionary.

Parameters:

• profile_dict (Dict[str, Dict[str, Tuple[int, ...]]]) –

  A dictionary mapping binding names to a dictionary containing min, opt, and max keys.

Returns:

• TensorRTProfile –

  A TensorRTProfile object.

Source code in model_navigator/api/config.py

@classmethod
def from_dict(cls, profile_dict: Dict[str, Dict[str, Tuple[int, ...]]]):
    """Create a TensorRTProfile from a dictionary.

    Args:
        profile_dict (Dict[str, Dict[str, Tuple[int, ...]]]):
            A dictionary mapping binding names to a dictionary containing ``min``, ``opt``, and
            ``max`` keys.

    Returns:
        TensorRTProfile:
            A TensorRTProfile object.
    """
    return cls({name: ShapeTuple(**shapes) for name, shapes in profile_dict.items()})

TensorType

Bases: Enum

All model formats supported by Model Navigator 'optimize' function.

TorchConfig dataclass

Bases: CustomConfigForFormat

Torch custom config used for TorchScript export.

Parameters:

• jit_type (Union[Union[str, JitType], Tuple[Union[str, JitType], ...]]) –

  Type of TorchScript export.

• strict (bool) –

  Enable or Disable strict flag for tracer used in TorchScript export, default: True.

format property

format: Format

Returns Format.TORCHSCRIPT.

Returns:

• Format –

  Format.TORCHSCRIPT

__post_init__

__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

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)
    self.strict = True

name classmethod

name()

Name of the config.

Source code in model_navigator/api/config.py

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

TorchTensorRTConfig dataclass

Bases: CustomConfigForFormat

Torch custom config used for TensorRT TorchScript conversion.

Parameters:

• precision (Union[Union[str, TensorRTPrecision], Tuple[Union[str, TensorRTPrecision], ...]]) –

  TensorRT precision.

• max_workspace_size (Optional[int]) –

  Max workspace size used by converter.

• trt_profile (Optional[TensorRTProfile]) –

  TensorRT profile.

format property

format: Format

Returns Format.TORCH_TRT.

Returns:

• Format –

  Format.TORCH_TRT

__post_init__

__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

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 classmethod

from_dict(config_dict)

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 config_dict.get("trt_profile") is not None 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()

Name of the config.

Source code in model_navigator/api/config.py

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

map_custom_configs

map_custom_configs(custom_configs)

Map custom configs from list to dictionary.

Parameters:

• custom_configs (Optional[Sequence[CustomConfig]]) –

  List of custom configs passed to API method

Returns:

• Dict –

  Mapped configs to dictionary

Source code in model_navigator/api/config.py

def map_custom_configs(custom_configs: Optional[Sequence[CustomConfig]]) -> Dict:
    """Map custom configs from list to dictionary.

    Args:
        custom_configs: List of custom configs passed to API method

    Returns:
        Mapped configs to dictionary
    """
    if not custom_configs:
        return {}

    return {config.name(): config for config in custom_configs}

model_navigator.api.MaxThroughputAndMinLatencyStrategy

Bases: RuntimeSearchStrategy

Get runtime with the highest throughput and the lowest latency.

model_navigator.api.MaxThroughputStrategy

Bases: RuntimeSearchStrategy

Get runtime with the highest throughput.

model_navigator.api.MinLatencyStrategy

Bases: RuntimeSearchStrategy

Get runtime with the lowest latency.