neural_compressor.compression.pruner.tf_criteria

Tensorflow pruning criterion.

Classes

PruningCriterion	Pruning base criterion.
MagnitudeCriterion	Pruning criterion.

Functions

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

Module Contents

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

Register a criterion to the registry.

neural_compressor.compression.pruner.tf_criteria.get_tf_criterion(config, modules)

Get registered criterion class.

class neural_compressor.compression.pruner.tf_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.

class neural_compressor.compression.pruner.tf_criteria.MagnitudeCriterion(modules, config)

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.