neural_compressor.experimental.pytorch_pruner.scheduler

scheduler module.

Module Contents

Classes

Scheduler

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.experimental.pytorch_pruner.scheduler.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. Define the scheduler type.

Returns:

The class of register.

Return type:

cls

neural_compressor.experimental.pytorch_pruner.scheduler.get_scheduler(config)

Get registered scheduler class.

Get a scheduler object from SCHEDULERS.

Parameters:

config – A config dict object. Contains the scheduler information.

Returns:

A Scheduler object.

class neural_compressor.experimental.pytorch_pruner.scheduler.Scheduler(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. Contains the scheduler information.

config

A config dict object. Contains the scheduler information.

abstract update_sparsity_ratio(aggressive_ratio, current_prune_step, total_prune_steps, masks)

To be implemented in subclasses.

class neural_compressor.experimental.pytorch_pruner.scheduler.OneshotScheduler(config)

Bases: Scheduler

Pruning Scheduler.

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

Parameters:

config – A config dict object. Contains the scheduler information.

Inherit from parent class Scheduler.
update_sparsity_ratio(aggressive_ratio, current_prune_step, total_prune_steps, masks)

Return the aggressive ratio.

class neural_compressor.experimental.pytorch_pruner.scheduler.IterativeScheduler(config)

Bases: Scheduler

Pruning Scheduler.

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

Parameters:

config – A config dict object. Contains the scheduler information.

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

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.

Returns：

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