neural_compressor.pruner.pruners

Pruner.

Module Contents

Classes

BasePruner

Pruning Pruner.

BasicPruner

Pruning Pruner.

PatternLockPruner

Pruning Pruner.

ProgressivePruner

Pruning Pruner.

Functions

register_pruner(name)

Class decorator to register a Pruner subclass to the registry.

parse_valid_pruner_types()

Get all valid pruner names.

get_pruner(config, modules)

Get registered pruner class.
neural_compressor.pruner.pruners.register_pruner(name)

Class decorator to register a Pruner subclass to the registry.

Decorator function used before a Pattern subclass. Make sure that the Pruner class decorated by this function can be registered in PRUNERS.

Parameters:

  • cls (class) – The subclass of register.

  • name – A string. Define the pruner type.

Returns:

The class of register.

Return type:

cls

neural_compressor.pruner.pruners.parse_valid_pruner_types()

Get all valid pruner names.

neural_compressor.pruner.pruners.get_pruner(config, modules)

Get registered pruner class.

Get a Pruner object from PRUNERS.

Parameters:

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

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

Returns:

A Pruner object.

Raises: AssertionError: Cuurently only support pruners that have been registered in PRUNERS.

class neural_compressor.pruner.pruners.BasePruner(config, modules)

Pruning Pruner.

The class which executes pruning process.

Parameters:

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

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

modules

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

config

A config dict object that contains the pruner information.

masks

A dict {“module_name”: Tensor} that stores the masks for modules’ weights.

scores

A dict {“module_name”: Tensor} that stores the score for modules’ weights, which are used to determine what parts to be pruned by a criterion.

pattern

A Pattern object defined in ./patterns.py

scheduler

A scheduler object defined in ./scheduler.py

current_sparsity_ratio

A float representing the current model’s sparsity ratio; it is initialized to be zero.

global_step

An integer representing the total steps the model has run.

start_step

An integer representing when to trigger pruning process.

end_step

An integer representing when to end pruning process.

pruning_frequency

An integer representing the pruning frequency; it is valid when iterative pruning is enabled.

target_sparsity_ratio

A float showing the final sparsity after pruning.

max_sparsity_ratio_per_op

A float showing the maximum sparsity ratio for every module.

on_epoch_begin(epoch)

Implement at the beginning of each epoch.

mask_weights()

Apply masks to corresponding modules’ weights.

Weights are multipled with masks. This is the formal pruning process.

mask_weights_general(input_masks)

Apply input masks to corresponding modules’ weights.

Weights are multipled with input_masks.

Parameters:

input_masks – A dict {“module_name”: Tensor} that stores the masks for modules’ weights.

on_step_begin(local_step)

Implement at the start of each step.

update_masks(local_step)

Update the masks at a given local step.

on_epoch_end()

Implement at the end of each epoch.

on_step_end()

Implement at the end of each step.

on_before_optimizer_step()

Implement before optimizer.step().

on_after_optimizer_step()

Implement after optimizer.step().

Prune the model after optimization.

on_train_begin()

Implement at the beginning of training phase.

on_train_end()

Implement at the end of training phase.

on_before_eval()

Implement at the beginning of evaluation phase.

on_after_eval()

Implement at the end of evaluation phase.

check_is_pruned_step(step)

Check if a pruning process should be performed at the current step.

Parameters:

step – an integer representing the number of current step.

Returns:

A Boolean.

class neural_compressor.pruner.pruners.BasicPruner(config, modules)

Bases: BasePruner

Pruning Pruner.

The class which executes pruning process. 1. Defines pruning functions called at step begin/end, epoch begin/end. 2. Defines the pruning criterion.

Parameters:

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

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

pattern

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

criterion

A Criterion Object that defines which weights are to be pruned

scheduler

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

reg

A Reg object that defines regulization terms.

set_global_step(global_step)

Set global step number.

update_masks(local_step)

Update the masks at a given local step.

on_before_optimizer_step()

Implement before optimizer.step().

on_after_optimizer_step()

Prune the model after optimization.

class neural_compressor.pruner.pruners.PatternLockPruner(config, modules)

Bases: BasePruner

Pruning Pruner.

A Pruner class derived from BasePruner. In this pruner, original model’s sparsity pattern will be fixed while training. This pruner is useful when a user trains a sparse model without changing its original structure.

Parameters:

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

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

Inherit from parent class Pruner.
update_masks(local_step)

Update the masks at a given local step.

on_after_optimizer_step()

Implement after optimizer.step().

Prune the model after optimization.

class neural_compressor.pruner.pruners.ProgressivePruner(config, modules)

Bases: BasicPruner

Pruning Pruner.

A Pruner class derived from BasePruner. In this pruner, mask interpolation will be applied. Mask interpolation is a fine-grained improvement for NxM structured pruning by adding interval

masks between masks of two pruning steps.

Parameters:

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

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

Inherit from parent class Pruner.
check_progressive_validity()

Check if the settings of progressive pruning are valid.

check_is_pruned_progressive_step(step)

Check if a progressive pruning process should be performed at the current step.

Parameters:

step – an integer representing the number of current step.

Returns:

A Boolean.

update_masks_progressive(local_step)

Update the masks in progressive pruning mode at a given local step.

on_step_begin(local_step)

Update the masks at a given local_step.

Implement at the start of each step.

on_before_optimizer_step()

Implement before optimizer.step().

on_after_optimizer_step()

Prune the model after optimization.

print_progressive_sparsity()

Output the progressive sparsity.