neural_compressor.pruner.schedulers

scheduler module.

Module Contents

Classes

PruningScheduler

Pruning Scheduler.

OneshotScheduler

Pruning Scheduler.

IterativeScheduler

Pruning Scheduler.

Functions

register_scheduler(name)

Class decorator used to register a Scheduler subclass to the registry.

get_scheduler(config)

Get registered scheduler class.
neural_compressor.pruner.schedulers.register_scheduler(name)

Class decorator used to register a Scheduler subclass to the registry.

Decorator function used before a Scheduler subclass. Make sure that the Scheduler class decorated by this function can be registered in SCHEDULERS.

Parameters:

  • cls (class) – The class of register.

  • name – A string that defines the scheduler type.

Returns:

The class of register.

Return type:

cls

neural_compressor.pruner.schedulers.get_scheduler(config)

Get registered scheduler class.

Get a scheduler object from SCHEDULERS.

Parameters:

config – A config dict object that contains the scheduler information.

Returns:

A Scheduler object.

class neural_compressor.pruner.schedulers.PruningScheduler(config)

Pruning Scheduler.

The class which defines a sparsity changing process during pruning. Mainly contains two types:

  1. iterative scheduler. Prune the model from dense to target sparsity gradually.

  2. one-shot scheduler. Prune the model in a single step and reach the target sparsity.

Parameters:

config – A config dict object that contains the scheduler information.

config

A config dict object that contains the scheduler information.

abstract update_sparsity_ratio(target_ratio, current_prune_step, total_prune_steps, masks, init_ratio=0.0)

To be implemented in subclasses.

class neural_compressor.pruner.schedulers.OneshotScheduler(config)

Bases: PruningScheduler

Pruning Scheduler.

A Scheduler class derived from Scheduler. Prune the model to target sparsity once.

Parameters:

config – A config dict object that contains the scheduler information.

Inherit from parent class Scheduler.
update_sparsity_ratio(target_ratio, current_prune_step, total_prune_steps, masks, init_ratio=0.0)

Update sparsity ratio.

Parameters:

  • target_ratio – A float representing the sparsity ratio after pruning.

  • current_prune_step – An integer representing the current pruning step.

  • total_prune_steps – An integer representing the total number of steps of the pruning process.

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

  • init_ratio – A float representing the sparsity ratio before pruning.

Returns:

A float representing the sparsity ratio that the model will reach after the next pruning step.

class neural_compressor.pruner.schedulers.IterativeScheduler(config)

Bases: PruningScheduler

Pruning Scheduler.

A Scheduler class derived from Scheduler. Prune the model from dense to target sparsity in several steps.

Parameters:

config – A config dict object that contains the scheduler information.

Inherit from parent class Scheduler.
update_sparsity_ratio(target_ratio, current_prune_step, total_prune_steps, masks, init_sparsity_ratio=0.0)

Obtain new target sparsity ratio according to the step.

Parameters:

  • target_ratio – A float. The target sparsity ratio.

  • current_prune_step – A integer. The current pruning step.

  • total_prune_steps – A integer. The total steps included in the pruning progress.

  • masks – A dict{“module_name”: Tensor}. The masks for modules’ weights.

  • init_sparsity_ratio –

Returns:

A float representing the target sparsity ratio the model will reach after the next pruning step.