bitorch.quantizations.quantization_scheduler.Quantization_Scheduler

class bitorch.quantizations.quantization_scheduler.Quantization_Scheduler(model: Module, steps: int, quantizations: List[Quantization], scheduling_procedure: str, schedule_all_quantizations: bool = False, exclude_layers: List[Type] = [])[source]

Methods

__init__

Initiates the quantization scheduler and replaces the activation function inside the model with scheduled quantizers

get_scheduled_quantizer

gets the scheduling class associated with the given scheduling procedure

replace_quantizations

replaces all quantization functions present in the model with a scheduled quantizer.

step

updates all instances of scheduled quantizers in the model

Attributes

procedure_classes

__init__(model: Module, steps: int, quantizations: List[Quantization], scheduling_procedure: str, schedule_all_quantizations: bool = False, exclude_layers: List[Type] = []) → None[source]

Initiates the quantization scheduler and replaces the activation function inside the model with scheduled quantizers

Parameters:

  • model (Module) – model to be scheduled quantized

  • steps (int) – number of steps, e.g. number of epochs. Each step the step() method has to be called once to update all scheduled quantizers.

  • quantizations (List[Quantization]) – Quantization functions to be scheduled

  • scheduling_procedure (str) – procedure to be used for scheduling. See available subclasses of ScheduledQuantizer

  • schedule_all_quantizations (bool) – toggles weather all quantizations in the model shall be replaced with quantized schedulers or weather only the quantized scheduler layers already present shall be used for scheduling. Defaults to False.

  • exclude_layers (List[Type], optional) – list of layers types to exclude from replacement with scheduled quantizers. Defaults to [].

get_scheduled_quantizer(procedure: str) → Type[source]

gets the scheduling class associated with the given scheduling procedure

Parameters:

procedure (str) – name of the scheduling procedure to be used

Returns:

a subclass of ScheduledQuantizer

Return type:

Type

replace_quantizations(model: Module, exclude_layers: List[Type], replace_all_quantizations: bool) → None[source]

replaces all quantization functions present in the model with a scheduled quantizer. iterates recursevely to the model layers.

Parameters:

  • model (Module) – model have the quantization functions replaced

  • exclude_layers (List[Type]) – list of layers to exclude from replacement, e.g. if QConv2d is specified, the quantization functions from all QConv2d layers (input and weight) are not replaced

  • replace_all_quantizations (bool) – toggles weather to replace all quantizations or just the instances of ScheduledQuantizer

step() → None[source]

updates all instances of scheduled quantizers in the model