neural_compressor.experimental.pruning

pruning module.

Module Contents

Classes

Pruning	This is base class of pruning object.
TfPruningCallback	Class that contains callback functions.

class neural_compressor.experimental.pruning.Pruning(conf_fname_or_obj=None)

Bases: neural_compressor.experimental.component.Component

This is base class of pruning object.

Since DL use cases vary in the accuracy metrics (Top-1, MAP, ROC etc.), loss criteria (<1% or <0.1% etc.) and pruning objectives (performance, memory footprint etc.). Pruning class provides a flexible configuration interface via YAML for users to specify these parameters.

Parameters:

conf_fname_or_obj (string or obj) – The path to the YAML configuration file or PruningConf class containing accuracy goal, pruning objective and related dataloaders etc.

conf

A config dict object. Contains pruning setting parameters.

pruners

A list of Pruner object.

property pruning_func

Not support get pruning_func.

property evaluation_distributed

Getter to know whether need distributed evaluation dataloader.

property train_distributed

Getter to know whether need distributed training dataloader.

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

Functions prepare for generate_hooks, generate_pruners.

pre_process()

Functions called before pruning begins, usually set up pruners.

execute()

Functions that execute the pruning process.

Follow the working flow: evaluate the dense model -> train/prune the model, evaluate the sparse model.

generate_hooks()

Register hooks for pruning.

generate_pruners()

Functions that generate pruners and set up self.pruners.

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.

class neural_compressor.experimental.pruning.TfPruningCallback(nc_model, input_model, hooks)

Bases: object

Class that contains callback functions.

Parameters:

• nc_model – A neural compression model object.

• hooks – A dict. Contains pure-defined hooks.

on_train_begin(logs=None, dataloader=None)

Call the same-name function from hooks.

on_train_end(logs=None)

Call the same-name function from hooks.

pre_epoch_begin(logs=None, dataloader=None)

Call the same-name function from hooks.

post_epoch_end(logs=None)

Call the same-name function from hooks.

on_epoch_begin(epoch, logs=None)

Call the same-name function from hooks.

on_epoch_end(logs=None)

Call the same-name function from hooks.

on_step_begin(batch, logs=None)

Call the same-name function from hooks.

on_batch_begin(batch, logs=None)

Call the same-name function from hooks.

on_after_compute_loss(input, s_outputs, s_loss, t_outputs=None)

Call the same-name function from hooks.

on_step_end(logs=None)

Call the same-name function from hooks.

on_batch_end(logs=None)

Call the same-name function from hooks.