neural_compressor.compression.pruner.pruning

Pruning.

Classes

BasePruning	Pruning.
BasicPruning	Pruning.
SparseGPTPruning	SparseGPT Pruning
RetrainFreePruning	Pruning.

Functions

register_pruning(name)

Class decorator to register a pruning subclass to the registry.

Module Contents

neural_compressor.compression.pruner.pruning.register_pruning(name)

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)

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

The model object to prune.

config_file_path

A string representing the path to a config file.

pruners

A list. A list of Pruner objects.

pruner_info

A config dict object that contains pruners’ information.

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

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

The model object to prune.

config_file_path

A string representing the path to a config file.

pruners

A list. A list of Pruner objects.

pruner_info

A config dict object that contains pruners’ information.

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

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')

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

The model object to prune.

config_file_path

A string representing the path to a config file.

pruners

A list. A list of Pruner objects.

pruner_info

A config dict object that contains pruners’ information.