neural_compressor.compression.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.compression.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.compression.pruner.pruners.parse_valid_pruner_types()

Get all valid pruner names.

neural_compressor.compression.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.compression.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.

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

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.

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

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.

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

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.