neural_compressor.compression.pruner

Pruning init.

Subpackages

Submodules

Package Contents

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.
neural_compressor.compression.pruner.save(obj: object, f, pickle_module=None, pickle_protocol=None, _use_new_zipfile_serialization=None)[source]

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

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.