Inplace Optimize

Model Navigator Inplace Optimize.

model_navigator.inplace.wrapper

Inplace Optimize model wrapper.

Module

Module(
    module,
    optimize_config=None,
    name=None,
    input_mapping=None,
    output_mapping=None,
    timer=None,
    offload_parameters_to_cpu=False,
)

Bases: ObjectProxy

Inplace Optimize module wrapper.

This class wraps a torch module and provides inplace optimization functionality. Depening on the configuration set in config, the module will be optimized, recorded, or passed through.

This wrapper can be used in place of a torch module, and will behave identically to the original module.

Parameters:

• module (Module) –

  torch module to wrap.

• optimize_config (Optional[OptimizeConfig], default: None ) –

  optimization configuration.

• name (Optional[str], default: None ) –

  module name.

• input_mapping (Optional[Callable], default: None ) –

  function to map module inputs to the expected input.

• output_mapping (Optional[Callable], default: None ) –

  function to map module outputs to the expected output.

• offload_parameters_to_cpu (bool, default: False ) –

  offload parameters to cpu.

Example

import torch import model_navigator as nav model = torch.nn.Linear(10, 10) model = nav.Module(model)

Initialize Module.

Source code in model_navigator/inplace/wrapper.py

def __init__(
    self,
    module: torch.nn.Module,
    optimize_config: Optional[OptimizeConfig] = None,
    name: Optional[str] = None,
    input_mapping: Optional[Callable] = None,
    output_mapping: Optional[Callable] = None,
    timer: Optional[Timer] = None,
    offload_parameters_to_cpu: bool = False,
) -> None:
    """Initialize Module."""
    super().__init__(module)
    self._optimize_config = optimize_config or OptimizeConfig()
    self._name = name or get_object_name(module)
    self._input_mapping = input_mapping or (lambda x: x)
    self._output_mapping = output_mapping or (lambda x: x)
    if timer:
        self._module_timer = timer.register_module(self._name)
    else:
        self._module_timer = None

    wrapper_cls = {
        Mode.OPTIMIZE: RecordAndOptimizeModule,
        Mode.RECORDING: RecordModule,
        Mode.RUN: OptimizedModule,
        Mode.PASSTHROUGH: PassthroughModule,
    }[inplace_config.mode]
    self._wrapper = wrapper_cls(
        module,
        self._optimize_config,
        self._name,
        self._input_mapping,
        self._output_mapping,
    )

    module_registry.register(self._name, self)

is_optimized property

is_optimized

Check if the module is optimized.

is_ready_for_optimization property

is_ready_for_optimization

Check if the module is ready for optimization.

name property

name

Module name.

optimize_config property

optimize_config

Module optimize config.

wrapper property

wrapper

Return the wrapper module.

__call__

__call__(*args, **kwargs)

Call the wrapped module.

This method overrides the call method of the wrapped module. If the module is already optimized it is replaced with the optimized one.

Source code in model_navigator/inplace/wrapper.py

def __call__(self, *args, **kwargs) -> Any:
    """Call the wrapped module.

    This method overrides the __call__ method of the wrapped module.
    If the module is already optimized it is replaced with the optimized one.
    """
    if self._module_timer and self._module_timer.enabled:
        with self._module_timer:
            output = self._wrapper(*args, **kwargs)
            if isinstance(self, torch.nn.Module):
                torch.cuda.synchronize()
    else:
        output = self._wrapper(*args, **kwargs)

    return output

load_optimized

load_optimized()

Load optimized module.

Source code in model_navigator/inplace/wrapper.py

def load_optimized(self) -> None:
    """Load optimized module."""
    assert self.is_optimized, f"Module {self.name} is not optimized."

    if not isinstance(self.wrapper, OptimizedModule):
        self._wrapper = OptimizedModule(
            self._wrapper._module,
            self._optimize_config,
            self._name,
            self._input_mapping,
            self._output_mapping,
        )

optimize

optimize()

Optimize the module.

Source code in model_navigator/inplace/wrapper.py

def optimize(self) -> None:
    """Optimize the module."""
    assert isinstance(self.wrapper, RecordModule), f"Module {self.name} must be in recording mode to optimize."
    assert not self.is_optimized, f"Module {self.name} is already optimized."
    assert hasattr(self.wrapper, "optimize"), f"Module {self.name} does not have an optimize method."
    self.wrapper.optimize()

module

module(module_callable=None, optimize_config=None, name=None, input_mapping=None, output_mapping=None)

Inplace Optimize module wrapper decorator.

This decorator wraps a torch module and provides inplace optimization functionality. Depening on the configuration set in config, the module will be optimized, recorded, or passed through.

This wrapper can be used in place of a torch module, and will behave identically to the original module.

Parameters:

• module_callable (Optional[Callable[[Any], Module]], default: None ) –

  decorated callable.

• optimize_config (Optional[OptimizeConfig], default: None ) –

  optimization configuration.

• name (Optional[str], default: None ) –

  module name.

• input_mapping (Optional[Callable], default: None ) –

  function to map module inputs to the expected input.

• output_mapping (Optional[Callable], default: None ) –

  function to map module outputs to the expected output.

Example

import torch import model_navigator as nav @nav.module ... def my_model(): ... return torch.nn.Linear(10, 10) model = my_model()

Source code in model_navigator/inplace/wrapper.py

def module(
    module_callable: Optional[Callable[[Any], torch.nn.Module]] = None,
    optimize_config: Optional[OptimizeConfig] = None,
    name: Optional[str] = None,
    input_mapping: Optional[Callable] = None,
    output_mapping: Optional[Callable] = None,
):
    """Inplace Optimize module wrapper decorator.

    This decorator wraps a torch module and provides inplace optimization functionality.
    Depening on the configuration set in config, the module will be
    optimized, recorded, or passed through.

    This wrapper can be used in place of a torch module, and will behave
    identically to the original module.

    Args:
        module_callable: decorated callable.
        optimize_config: optimization configuration.
        name: module name.
        input_mapping: function to map module inputs to the expected input.
        output_mapping: function to map module outputs to the expected output.

    Example:
        >>> import torch
        >>> import model_navigator as nav
        >>> @nav.module
        ... def my_model():
        ...     return torch.nn.Linear(10, 10)
        >>> model = my_model()
    """
    if module_callable is None:
        return functools.partial(
            module,
            optimize_config=optimize_config,
            name=name,
            input_mapping=input_mapping,
            output_mapping=output_mapping,
        )

    @wrapt.decorator
    def wrap_module(wrapped, instance, args, kwargs):
        return Module(
            wrapped(*args, **kwargs),
            optimize_config=optimize_config,
            name=name,
            input_mapping=input_mapping,
            output_mapping=output_mapping,
        )

    return wrap_module(module_callable)

model_navigator.inplace.config

Inplace Optimize configuration.

InplaceConfig

InplaceConfig()

Inplace Optimize configuration.

Initialize InplaceConfig.

Source code in model_navigator/inplace/config.py

def __init__(self) -> None:
    """Initialize InplaceConfig."""
    self._mode: Mode = DEFAULT_MODE
    self._cache_dir: pathlib.Path = DEFAULT_CACHE_DIR
    self._min_num_samples: int = DEFAULT_MIN_NUM_SAMPLES
    self._max_num_samples_stored: int = DEFAULT_MAX_NUM_SAMPLES_STORED
    self.strategy: RuntimeSearchStrategy = MinLatencyStrategy()

cache_dir property writable

cache_dir

Get the cache directory.

max_num_samples_stored property writable

max_num_samples_stored

Get the minimum number of samples to collect before optimizing.

min_num_samples property writable

min_num_samples

Get the minimum number of samples to collect before optimizing.

mode property writable

mode

Get the mode of the inplace Optimize.

Mode

Bases: Enum

Mode of the inplace Optimize.

OPTIMIZE: record registered models and optimize them when enough samples are collected. RECORDING: record registered models. RUN: replace registered models with optimized ones. PASSTHROUGH: do nothing.

OptimizeConfig dataclass

Configuration for inplace Optimize.

Parameters:

• sample_count (int, default: DEFAULT_SAMPLE_COUNT ) –

  Limits how many samples will be used from dataloader

• batching (Optional[bool], default: True ) –

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

• input_names (Optional[Tuple[str, ...]], default: None ) –

  Model input names

• output_names (Optional[Tuple[str, ...]], default: None ) –

  Model output names

• target_formats (Optional[Tuple[Union[str, Format], ...]], default: None ) –

  Target model formats for optimize process

• target_device (Optional[DeviceKind], default: CUDA ) –

  Target device for optimize process, default is CUDA

• runners (Optional[Tuple[Union[str, Type[NavigatorRunner]], ...]], default: None ) –

  Use only runners provided as parameter

• optimization_profile (Optional[OptimizationProfile], default: None ) –

  Optimization profile for conversion and profiling

• workspace (Optional[Path], default: None ) –

  Workspace where packages will be extracted

• verbose (Optional[bool], default: False ) –

  Enable verbose logging

• debug (Optional[bool], default: False ) –

  Enable debug logging from commands

• verify_func (Optional[VerifyFunction], default: None ) –

  Function for additional model verification

• custom_configs (Optional[Sequence[CustomConfig]], default: None ) –

  Sequence of CustomConfigs used to control produced artifacts

to_dict

to_dict()

Convert OptimizeConfig to dictionary.

Source code in model_navigator/inplace/config.py

def to_dict(self) -> Dict[str, Any]:
    """Convert OptimizeConfig to dictionary."""
    config_dict = {}
    for field in dataclasses.fields(self):
        value = getattr(self, field.name)
        config_dict[field.name] = value
    return config_dict

model_navigator.inplace

Inplace Optimize.

optimize

optimize()

Optimize all registered modules.

Source code in model_navigator/inplace/__init__.py

def optimize() -> None:
    """Optimize all registered modules."""
    module_registry.optimize()