mode — Model Optimizer 0.31.0 (original) (raw)

TensorRT Model Optimizer

This module contains the mode descriptor for the quantization mode.

Classes

class BaseCalibrateModeDescriptor

Bases: ModeDescriptor

Base class for quantization calibration algorithm modes.

All calibration algorithm modes must be derived from this base class. In addition, the config_class for the mode must return a subclass of QuantizeAlgorithmConfig.

This base class also provides some convenient wrappers/utilities for calibration algorithms to be translated into ModelOpt mode.

It includes:

1. A utility to convert the algorithm name to a mode name. This is useful since many algorithm names
   are trivial and not a good fit as a mode name. For example, "max" or None.
2. Conversion of the algorithm and kwargs arguments of
   calibrate API to a mode config list compatible with apply_mode.
3. Wrapper for the calibration functions in modelopt.torch.quantization.model_calib to be
   compatible with the ModelOpt convert entrypoint.

__init__(*args, **kwargs)

Initialize Base calibrate mode descriptor.

abstract property config_class_: type[QuantizeAlgorithmConfig]_

Specifies the config class for the mode.

property convert_: Callable[[Module, ModeloptBaseConfig], tuple[Module, dict[str, Any]]] | Callable[[Module, ModeloptBaseConfig, Any], tuple[Module, dict[str, Any]]]_

The calibrate algorithm mode’s entrypoint for converting a model.

This method is called by the ModelOpt framework when applying this calibration mode to a model. See wrapped_calib_func() for more details on the logic.

Note: Subclasses must specify the _calib_func class attribute with the appropriate calibration function to be used or override this method.

property name_: str_

Returns the value (str representation) of the mode.

property restore_: Callable[[Module, ModeloptBaseConfig, dict[str, Any]], Module]_

The mode’s entrypoint for restoring a model.

property update_for_new_mode_: Callable[[Module, ModeloptBaseConfig, dict[str, Any]], None]_

The mode’s entrypoint for updating the models state before new mode.

property update_for_save_: Callable[[Module, ModeloptBaseConfig, dict[str, Any]], None]_

The mode’s entrypoint for updating the models state before saving.