neural_compressor.compression.pruner

Pruning init.

Subpackages

• neural_compressor.compression.pruner.model_slim
• neural_compressor.compression.pruner.patterns
• neural_compressor.compression.pruner.pruners
• neural_compressor.compression.pruner.wanda

Submodules

• neural_compressor.compression.pruner.criteria
• neural_compressor.compression.pruner.pruning
• neural_compressor.compression.pruner.regs
• neural_compressor.compression.pruner.schedulers
• neural_compressor.compression.pruner.tf_criteria
• neural_compressor.compression.pruner.utils

Functions

save(obj, f[, pickle_module, pickle_protocol, ...])	A rewrite function for torch save.
prepare_pruning(model, config[, optimizer, ...])	Get registered pruning class, wrapper the model and optimizer to support all the pruning functionality.

Package Contents

neural_compressor.compression.pruner.save(obj: object, f, pickle_module=None, pickle_protocol=None, _use_new_zipfile_serialization=None)

A rewrite function for torch save.

Parameters:

• obj

• f

• pickle_module

• pickle_protocol

• _use_new_zipfile_serialization

Returns:

neural_compressor.compression.pruner.prepare_pruning(model, config, optimizer=None, dataloader=None, loss_func=None, framework='pytorch', device: str = None)

Get registered pruning class, wrapper the model and optimizer to support all the pruning functionality.

Get a pruning object from PRUNINGS.

Parameters:

• modules – A dict {“module_name”: Tensor} that stores the pruning modules’ weights.

• config – A config dict object that contains the pruners information.

Returns:

A pruning object.

Raises: AssertionError: Currently only support prunings that have been registered in PRUNINGS.