neural_compressor.compression.pruner.pruning

Pruning.

Module Contents

Classes

BasePruning

Pruning.

BasicPruning

Pruning.

SparseGPTPruning

SparseGPT Pruning

RetrainFreePruning

Pruning.

Functions

register_pruning(name)

Class decorator to register a pruning subclass to the registry.

neural_compressor.compression.pruner.pruning.register_pruning(name)[source]

Class decorator to register a pruning subclass to the registry.

Decorator function used before a pruner subclass. Make sure that the pruning class decorated by this function can be registered in PRUNINGS.

Parameters:
  • cls (class) – The subclass of register.

  • name – A string. Define the pruning type.

Returns:

The class of register.

Return type:

cls

class neural_compressor.compression.pruner.pruning.BasePruning(config, model, opt=None)[source]

Pruning.

The main class to do pruning; it contains at least one Pruner object.

Parameters:

config – a string representing the path to a config file. For config file template, please refer to https://github.com/intel/neural-compressor/tree/master/examples/pytorch/nlp/huggingface_models/text-classification/pruning/pytorch_pruner/eager/

model[source]

The model object to prune.

config_file_path[source]

A string representing the path to a config file.

pruners[source]

A list. A list of Pruner objects.

pruner_info[source]

A config dict object that contains pruners’ information.

class neural_compressor.compression.pruner.pruning.BasicPruning(config, model, opt=None)[source]

Pruning.

The main class to do pruning; it contains at least one Pruner object.

Parameters:

config – a string representing the path to a config file. For config file template, please refer to https://github.com/intel/neural-compressor/tree/master/examples/pytorch/nlp/huggingface_models/text-classification/pruning/pytorch_pruner/eager/

model[source]

The model object to prune.

config_file_path[source]

A string representing the path to a config file.

pruners[source]

A list. A list of Pruner objects.

pruner_info[source]

A config dict object that contains pruners’ information.

class neural_compressor.compression.pruner.pruning.SparseGPTPruning(config, model, dataloader, framework='pytorch', device=None)[source]

SparseGPT Pruning The SparseGPT pruning_class is derived from BasePruning.

Parameters:
  • config – A config dict object that contains the pruner information.

  • model – The model that need to be pruned.

  • dataloader – Processed datasets, which is necessary for sparseGPT pruning.

  • device – available device of pruning.

class neural_compressor.compression.pruner.pruning.RetrainFreePruning(config, model, dataloader=None, loss_func=None, framework='pytorch')[source]

Pruning.

The main class to do pruning; it contains at least one Pruner object.

Parameters:

config – a string representing the path to a config file. For config file template, please refer to https://github.com/intel/neural-compressor/tree/master/examples/pytorch/nlp/huggingface_models/text-classification/pruning/pytorch_pruner/eager/

model[source]

The model object to prune.

config_file_path[source]

A string representing the path to a config file.

pruners[source]

A list. A list of Pruner objects.

pruner_info[source]

A config dict object that contains pruners’ information.