neural_compressor.pruner.pruner_legacy.magnitude

Magnitude pruner.

Module Contents

Classes

BasicMagnitudePruner

Magnitude pruner class.

class neural_compressor.pruner.pruner_legacy.magnitude.BasicMagnitudePruner(model, local_config, global_config)

Bases: neural_compressor.pruner.pruner_legacy.pruner.Pruner

Magnitude pruner class.

Parameters:

• model (object) – The original model (currently PyTorchModel instance).

• local_config (Conf) – configs specific for this pruning instance.

• global_config (Conf) – global configs which may be overwritten by local_config.

on_epoch_begin(epoch)

Update target sparsity according to the schedule and compute mask accordingly.

on_step_begin(batch_id)

Apply mask to the weight.

compute_mask()

Compute masks according to absolute values.

on_epoch_end()

Sparsity ratio summary and apply mask to the weight.

on_step_end()

Apply mask to the weight.