neural_compressor.compression.pruner.criteria

Pruning criterion.

Classes

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

Functions

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

Module Contents

neural_compressor.compression.pruner.criteria.register_criterion(name)

Register a criterion to the registry.

neural_compressor.compression.pruner.criteria.get_criterion(config, modules, pattern, masks=None)

Get registered criterion class.

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

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.

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

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.

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

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.

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

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.

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

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.

class neural_compressor.compression.pruner.criteria.BlockMaskCriterion(modules, config, pattern, masks, alpha=0.9, beta=1.0)

Pruning criterion.

The block_mask criterion_class is derived from PruningCriterion. A momentum mechanism is used to calculate snip score, which determines if a block of weights 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.

class neural_compressor.compression.pruner.criteria.RetrainFreeCriterion(modules, config, pattern, masks)

Pruning criterion.

The retrain_free criterion_class is derived from PruningCriterion.

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.