:py:mod:`neural_compressor.training`
====================================

.. py:module:: neural_compressor.training

.. autoapi-nested-parse::

   The configuration of the training loop.



Module Contents
---------------

Classes
~~~~~~~

.. autoapisummary::

   neural_compressor.training.CompressionManager
   neural_compressor.training.CallBacks



Functions
~~~~~~~~~

.. autoapisummary::

   neural_compressor.training.fit
   neural_compressor.training.prepare_compression



.. py:class:: CompressionManager(model, callbacks_list)


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

   :param model: A model to be compressed. It should be neural compressor model.
   :param 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")


.. py:function:: fit(compression_manager, train_func, eval_func=None, eval_dataloader=None, eval_metric=None, **kwargs)

   Compress the model with tuning for quantization.

   :param compression_manager: The Compression manager contains the model and
                               callbacks.
   :type compression_manager: CompressionManager
   :param train_func: 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.
   :type train_func: function, optional
   :param eval_func: 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
   :type eval_func: function, optional
   :param eval_dataloader: 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.
   :type eval_dataloader: generator, optional
   :param eval_metric: Set metric class or a dict of built-in metric configures,
                       and neural_compressor will initialize this class when evaluation.
   :type eval_metric: dict or obj


.. py:function:: prepare_compression(model: Callable, confs: Union[Callable, List], **kwargs)

   Summary.

   :param model: The model to optimize.
   :type model: Callable, optional
   :param confs: Config of Distillation, Quantization, Pruning,
                 or list of config for orchestration optimization.
                 The config class is QuantizationAwareTrainingConfig,
                 PruningConfig, distillationConfig.
   :type confs: Union[Callable, List]
   :param options: The configure for random_seed, workspace,
                   resume path and tensorboard flag.
   :type options: Options, optional

   :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()


.. py:class:: CallBacks(callbacks_list)


   Define the basic command for the training loop.