neural_compressor.compression.pruner.pruners.retrain_free

Retrain free pruner.

Module Contents

Classes

PytorchRetrainFreePruner

Pruning Pruner.
class neural_compressor.compression.pruner.pruners.retrain_free.PytorchRetrainFreePruner(config, modules)[source]

Pruning Pruner. The retrain_free pruner_class is derived from BasePruner. This pruner references the mask search and mask rearrangement strategies in fast retraining free. RetrainFreePruner supports one-shot pruning (same effect as fast retraining free) and iterative pruning. Please refer to A Fast Post-Training Pruning Framework for Transformers

  1. Defines pruning functions called at step begin/end, before/after optimize and epoch begin/end.

  2. Defines the pruning criterion and fixed weight parameters.

  3. Obtain block masks and its grads.

  4. Rearrange block masks.

Parameters:

  • modules – A dict {“module_name”: Tensor} that stores the pruning modules’ weights.

  • config – A config dict object that contains the pruner information.

pattern[source]

A Pattern object that defines pruning weights’ arrangements within space.

criterion[source]

A Criterion Object that defines which weights are to be pruned

scheduler[source]

A Scheduler object that defines how the model’s sparsity changes as training/pruning proceeds.

reg[source]

A Reg object that defines regulization terms.