neural_compressor.mix_precision

Mix Precision for Neural Compressor.

Functions

fit(model, conf[, eval_func, eval_dataloader, eval_metric])

Fit low precision model generation across multiple framework backends.

Module Contents

neural_compressor.mix_precision.fit(model, conf, eval_func=None, eval_dataloader=None, eval_metric=None, **kwargs)[source]

Fit low precision model generation across multiple framework backends.

Parameters:

  • model (object) – For Tensorflow model, it could be a path to frozen pb, loaded graph_def object or a path to ckpt/savedmodel folder. For PyTorch model, it’s torch.nn.model instance. For onnx model, it should be a path to .onnx file or onnx.onnx_ml_pb2.ModelProto. For MXNet model, it’s mxnet.symbol.Symbol or gluon.HybirdBlock instance.

  • conf (MixedPrecisionConfig) – The MixedPrecisionConfig class containing accuracy goal, tuning objective and mixed_precision tuning space etc.

  • eval_func (function, optional) – The evaluation function provided by user. This function takes model as parameter, and evaluation dataset and metrics should be encapsulated in this function implementation and outputs a higher-is-better accuracy scalar value.

  • eval_dataloader (generator, optional) – Data loader for evaluation. It is iterable and should yield a tuple of (input, label). The input could be a object, list, tuple or dict, depending on user implementation, as well as it can be taken as model input. The label should be able to take as input of supported metrics. If this parameter is not None, user needs to specify pre-defined evaluation metrics through configuration file and should set “eval_func” parameter as None. Tuner will combine model, eval_dataloader and pre-defined metrics to run evaluation process.

  • eval_metric (obj, optional) – An Accuracy object that measures metric for mixed_precision.

Returns:

A Mixed precision model across various DL frameworks.

Raises:

AssertionError. –

Example:

from neural_compressor import mix_precision
from neural_compressor.config import MixedPrecisionConfig

conf = MixedPrecisionConfig()
converted_model = mix_precision.fit(model, conf=conf)