neural_compressor.experimental.compression.pruning

Pruning.

Module Contents

Functions

save(obj, f[, pickle_module, pickle_protocol, ...])

A rewrite function for torch save.

prepare_pruning(config, model, opt)

Wrapper the model and optimizer to support all the pruning functionality.
neural_compressor.experimental.compression.pruning.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.experimental.compression.pruning.prepare_pruning(config, model: torch.nn.Module, opt: torch.optim)[source]

Wrapper the model and optimizer to support all the pruning functionality.

Parameters:

  • config – WeightPruningConfig

  • model – The user’s model

  • opt – The user’s optimizer

Returns:

The modified model and optimizer