neural_compressor.pruner.criteria

pruning criterion.

Module Contents

Classes

PruningCriterion	Pruning base criterion.
MagnitudeCriterion	Pruning criterion.
GradientCriterion	Pruning criterion.
SnipCriterion	Pruning criterion.
SnipMomentumCriterion	Pruning criterion.

Functions

register_criterion(name)	Register a criterion to the registry.
get_criterion(config, modules)	Get registered criterion class.

neural_compressor.pruner.criteria.register_criterion(name)

Register a criterion to the registry.

neural_compressor.pruner.criteria.get_criterion(config, modules)

Get registered criterion class.

class neural_compressor.pruner.criteria.PruningCriterion(modules, config)

Pruning base criterion.

Parameters:

• config – A config dict object that includes information about pruner and pruning criterion.

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

scores

A dict {“module_name”: Tensor} that stores the scores of pruning modules.

on_step_begin()

Calculate and store the pruning scores of pruning modules at the beginning of a step.

on_before_optimizer_step()

Calculate and store the pruning scores of pruning modules before the optimizer step.

on_after_optimizer_step()

Calculate and store the pruning scores of pruning modules after the optimizer step.

class neural_compressor.pruner.criteria.MagnitudeCriterion(modules, config)

Bases: PruningCriterion

Pruning criterion.

The magnitude criterion_class is derived from PruningCriterion. The magnitude value is used to score and determine if a weight is to be pruned.

Parameters:

• config – A config dict object that includes information about pruner and pruning criterion.

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

scores

A dict {“module_name”: Tensor} that stores the scores of pruning modules.

on_step_begin()

Calculate and store the pruning scores based on magtinude criterion.

class neural_compressor.pruner.criteria.GradientCriterion(modules, config)

Bases: PruningCriterion

Pruning criterion.

The gradient criterion_class is derived from PruningCriterion. The absolute value of gradient is used to score and determine if a weight is to be pruned.

Parameters:

• config – A config dict object that includes information about pruner and pruning criterion.

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

scores

A dict {“module_name”: Tensor} that stores the scores of pruning modules.

on_before_optimizer_step()

Calculate and store the pruning scores based on gradient criterion.

class neural_compressor.pruner.criteria.SnipCriterion(modules, config)

Bases: PruningCriterion

Pruning criterion.

The snip criterion_class is derived from PruningCriterion. The product of magnitude and gradient is used to score and determine if a weight is to be pruned. Please refer to SNIP: Single-shot Network Pruning based on Connection Sensitivity. (https://arxiv.org/abs/1810.02340)

Parameters:

• config – A config dict object that includes information about pruner and pruning criterion.

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

scores

A dict {“module_name”: Tensor} that stores the scores of pruning modules.

on_before_optimizer_step()

Calculate and store the pruning scores based on snip criterion.

class neural_compressor.pruner.criteria.SnipMomentumCriterion(modules, config)

Bases: PruningCriterion

Pruning criterion.

The snip_momentum criterion_class is derived from PruningCriterion. A momentum mechanism is used to calculate snip score, which determines if a weight is to be pruned.

Parameters:

• config – A config dict object that includes information about pruner and pruning criterion.

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

• alpha – A parameter that determines how much of the snip score is preserved from last pruning step.

• beta – A parameter that determines how much of the snip score is updated at the current step.

scores

A dict {“module_name”: Tensor} that stores the scores of pruning modules.

on_before_optimizer_step()

Calculate and store the pruning scores based on snip_momentum criterion.