Shortcuts

QuantizationAwareTraining

class pytorch_lightning.callbacks.QuantizationAwareTraining(qconfig='fbgemm', observer_type='average', collect_quantization=None, modules_to_fuse=None, input_compatible=True)[source]

Bases: pytorch_lightning.callbacks.base.Callback

Quantization allows speeding up inference and decreasing memory requirements by performing computations and storing tensors at lower bitwidths (such as INT8 or FLOAT16) than floating point precision. We use native PyTorch API so for more information see Quantization.

Warning

QuantizationAwareTraining is in beta and subject to change.

Parameters
  • qconfig (Union[str, QConfig]) –

    quantization configuration:

  • observer_type (str) – allows switching between MovingAverageMinMaxObserver as “average” (default) and HistogramObserver as “histogram” which is more computationally expensive.

  • collect_quantization (Union[Callable, int, None]) –

    count or custom function to collect quantization statistics:

    • None (deafult). The quantization observer is called in each module forward

      (useful for collecting extended statistic when useing image/data augmentation).

    • int. Use to set a fixed number of calls, starting from the beginning.

    • Callable. Custom function with single trainer argument.

      See this example to trigger only the last epoch:

      def custom_trigger_last(trainer):
          return trainer.current_epoch == (trainer.max_epochs - 1)
      
      QuantizationAwareTraining(collect_quantization=custom_trigger_last)
      

  • modules_to_fuse (Optional[Sequence]) – allows you fuse a few layers together as shown in diagram to find which layer types can be fused, check https://github.com/pytorch/pytorch/pull/43286.

  • input_compatible (bool) – preserve quant/dequant layers. This allows to feat any input as to the original model, but break compatibility to torchscript.

on_fit_end(trainer, pl_module)[source]

Called when fit ends

on_fit_start(trainer, pl_module)[source]

Called when fit begins

Read the Docs v: stable
Versions
latest
stable
1.2.2
1.2.1
1.2.0
1.1.8
1.1.7
1.1.6
1.1.5
1.1.4
1.1.3
1.1.2
1.1.1
1.1.0
1.0.8
1.0.7
1.0.6
1.0.5
1.0.4
1.0.3
1.0.2
1.0.1
1.0.0
0.10.0
0.9.0
0.8.5
0.8.4
0.8.3
0.8.2
0.8.1
0.8.0
0.7.6
0.7.5
0.7.4
0.7.3
0.7.2
0.7.1
0.7.0
0.6.0
0.5.3.2
0.5.3
0.4.9
release-1.2-dev
release-1.0.x
Downloads
pdf
html
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.