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

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

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

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

A config dict object. Contains the scheduler information.

class neural_compressor.experimental.pytorch_pruner.scheduler.OneshotScheduler(config)[source]

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.
class neural_compressor.experimental.pytorch_pruner.scheduler.IterativeScheduler(config)[source]

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.