neural_compressor.experimental.pytorch_pruner.pruning

pruning module.

Module Contents

Classes

Pruning

Pruning.

class neural_compressor.experimental.pytorch_pruner.pruning.Pruning(config)

Pruning.

The main class that users will used in codes to do pruning. Contain at least one Pruner object.

Parameters:

config – a string. The path to a config file. For config file template, please refer to https://github.com/intel/neural-compressor/tree/master/examples/pytorch/nlp/huggingface_models/text-classification/pruning/pytorch_pruner/eager/

model

The model object to prune.

config_file_path

A string. The path to a config file.

pruners

A list. A list of Pruner objects.

pruner_info

A config dict object. Contains pruners’ information.

update_items_for_all_pruners(**kwargs)

Functions which add User-defined arguments to the original configurations.

The original config of pruning is read from a file. However, users can still modify configurations by passing key-value arguments in this function. Please note that the key-value arguments’ keys are analysable in current configuration.

prepare()

Align with old API’s calling pipeline.

get_sparsity_ratio()

Functions that calculate a modules/layers sparsity.

Returns:

Three floats. elementwise_over_matmul_gemm_conv refers to zero elements’ ratio in pruning layers. elementwise_over_all refers to zero elements’ ratio in all layers in the model. blockwise_over_matmul_gemm_conv refers to all-zero blocks’ ratio in pruning layers.

on_train_begin()

Functions called in the beginning of training process.

Before training, ensure that pruners are generated.

on_epoch_begin(epoch)

Functions called in the beginning of every epoch.

on_step_begin(local_step)

Functions called in the beginning of every step.

on_before_optimizer_step()

Functions called before optimizer.step().

on_step_end()

Functions called in the end of every step.

on_epoch_end()

Functions called in the end of every epoch.

on_train_end()

Functions called in the end of training.

on_before_eval()

Functions called in the beginning of evaluation.

on_after_eval()

Functions called in the end of evaluation.

on_after_optimizer_step()

Functions called after optimizer.step().