neural_compressor.training

The configuration of the training loop.

Module Contents

Classes

CompressionManager	CompressionManager is uesd in train loop for what user want to deal with additional.
CallBacks	Define the basic command for the training loop.

Functions

fit(compression_manager, train_func[, eval_func, ...])	Compress the model with tuning for quantization.
prepare_compression(model, confs, **kwargs)	Summary.

class neural_compressor.training.CompressionManager(model, callbacks_list)

CompressionManager is uesd in train loop for what user want to deal with additional.

Parameters:

• model – A model to be compressed. It should be neural compressor model.

• callbacks – A list of Callbacks instances. Such as: DistillationCallbbacks, QuantizationAwareTrainingCallbacks, PruningCallbacks.

Examples:

import neural_compressor.training.prepare_compression
compression_manager = prepare_compression(nc_model, confs)
compression_manager.callbacks.on_train_begin()
model = compression_manager.model
train_loop:
    for epoch in range(epochs):
        compression_manager.callbacks.on_epoch_begin(epoch)
        for i, batch in enumerate(dataloader):
            compression_manager.callbacks.on_step_begin(i)
            ......
            output = model(batch)
            loss = ......
            loss = compression_manager.callbacks.on_after_compute_loss(batch, output, loss)
            loss.backward()
            compression_manager.callbacks.on_before_optimizer_step()
            optimizer.step()
            compression_manager.callbacks.on_step_end()
        compression_manager.callbacks.on_epoch_end()
compression_manager.callbacks.on_train_end()
compression_manager.save("path_to_save")

neural_compressor.training.fit(compression_manager, train_func, eval_func=None, eval_dataloader=None, eval_metric=None, **kwargs)

Compress the model with tuning for quantization.

Parameters:

• compression_manager (CompressionManager) – The Compression manager contains the model and callbacks.

• train_func (function, optional) – Training function for quantization aware training. It is optional. This function takes “model” as input parameter and executes entire inference process. If this parameter specified.

• 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. The pseudo code should be something like: def eval_func(model):

  input, label = dataloader() output = model(input) accuracy = metric(output, label) return accuracy

• 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 object and should set “eval_func” paramter as None. Tuner will combine model, eval_dataloader and pre-defined metrics to run evaluation process.

• eval_metric (dict or obj) – Set metric class or a dict of built-in metric configures, and neural_compressor will initialize this class when evaluation.

neural_compressor.training.prepare_compression(model: Callable, confs: Callable | List, **kwargs)

Summary.

Parameters:

• model (Callable, optional) – The model to optimize.

• confs (Union[Callable, List]) – Config of Distillation, Quantization, Pruning, or list of config for orchestration optimization. The config class is QuantizationAwareTrainingConfig, PruningConfig, distillationConfig.

• options (Options, optional) – The configure for random_seed, workspace, resume path and tensorboard flag.

Returns:

CompressionManager

Examples:

import neural_compressor.training.prepare_compression

compression_manager = prepare_compression(conf, model)
train_loop:
    compression_manager.on_train_begin()
    for epoch in range(epochs):
        compression_manager.on_epoch_begin(epoch)
        for i, batch in enumerate(dataloader):
            compression_manager.on_step_begin(i)
            ......
            output = model(batch)
            loss = ......
            loss = compression_manager.on_after_compute_loss(batch, output, loss)
            loss.backward()
            compression_manager.on_before_optimizer_step()
            optimizer.step()
            compression_manager.on_step_end()
        compression_manager.on_epoch_end()
    compression_manager.on_train_end()

class neural_compressor.training.CallBacks(callbacks_list)

Define the basic command for the training loop.