Skip to content

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)

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.

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 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
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,
) -> 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)

    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._wrapper)

name property 

name: str

Module name.

optimize_config property 

optimize_config: OptimizeConfig

Module optimize config.

__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 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
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 isinstance(self._wrapper, RecordModule) and self._wrapper._optimized:
        self._wrapper = OptimizedModule(
            self._wrapper._module,
            self._optimize_config,
            self._name,
            self._input_mapping,
            self._output_mapping,
        )
    output = self._wrapper(*args, **kwargs)
    return output

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 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
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 
50
51
52
53
54
55
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.strategy: RuntimeSearchStrategy = MinLatencyStrategy()

cache_dir property writable 

cache_dir: pathlib.Path

Get the cache directory.

min_num_samples property writable 

min_num_samples: int

Get the minimum number of samples to collect before optimizing.

mode property writable 

mode: 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[pathlib.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 
124
125
126
127
128
129
130
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 
19
20
21
def optimize() -> None:
    """Optimize all registered modules."""
    module_registry.optimize()