neural_compressor.compression.pruner.criteria

Pruning criterion.

Module Contents

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.
neural_compressor.compression.pruner.criteria.register_criterion(name)[source]

Register a criterion to the registry.

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

Get registered criterion class.

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

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[source]

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

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

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[source]

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

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

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[source]

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

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

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[source]

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

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

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[source]

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)[source]

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[source]

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

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

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[source]

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