neural_compressor.experimental.component

This is a module for Component class.

The Component class will be inherited by the class ‘Quantization’, ‘Pruning’ and ‘Distillation’.

Module Contents

Classes

Component

This is base class of Neural Compressor Component.

class neural_compressor.experimental.component.Component(conf_fname_or_obj=None, combination=None)

Bases: object

This is base class of Neural Compressor Component.

This class will be inherited by the class ‘Quantization’, ‘Pruning’ and ‘Distillation’. This design is mainly for one-shot optimization for pruning/distillation/quantization-aware training. In this class will apply all hooks for ‘Quantization’, ‘Pruning’ and ‘Distillation’.

property train_func

Not support get train_func.

property eval_func

Not support get eval_func.

property train_dataloader

Getter to train dataloader.

property eval_dataloader

Getter to eval dataloader.

property model

Getter of model in neural_compressor.model.

prepare()

Register Quantization Aware Training hooks.

prepare_qat()

Register Quantization Aware Training hooks.

pre_process()

Initialize some attributes, such as the adaptor, the dataloader and train/eval functions from yaml config.

Component base class provides default function to initialize dataloaders and functions from user config. And for derived classes(Pruning, Quantization, etc.), an override function is required.

execute()

Execute the processing of this compressor.

Component base class provides compressing processing. And for derived classes(Pruning, Quantization, etc.), an override function is required.

post_process()

Post process after execution.

For derived classes(Pruning, Quantization, etc.), an override function is required.

on_train_begin(dataloader=None)

Be called before the beginning of epochs.

on_train_end()

Be called after the end of epochs.

pre_epoch_begin(dataloader=None)

Be called before the beginning of epochs.

post_epoch_end()

Be called after the end of epochs.

on_epoch_begin(epoch)

Be called on the beginning of epochs.

on_step_begin(batch_id)

Be called on the beginning of batches.

on_batch_begin(batch_id)

Be called on the beginning of batches.

on_after_compute_loss(input, student_output, student_loss, teacher_output=None)

Be called on the end of loss computation.

on_before_optimizer_step()

Be called before optimizer step.

on_after_optimizer_step()

Be called after optimizer step.

on_before_eval()

Be called before evaluation.

on_after_eval()

Be called after evaluation.

on_post_grad()

Be called before optimizer step.

on_step_end()

Be called on the end of batches.

on_batch_end()

Be called on the end of batches.

on_epoch_end()

Be called on the end of epochs.

register_hook(scope, hook, input_args=None, input_kwargs=None)

Register hook for component.

Input_args and input_kwargs are reserved for user registered hooks.