diff --git a/model_compression_toolkit/qat/keras/README.md b/model_compression_toolkit/qat/keras/README.md index 8b3fe81fe..f1cf6593e 100644 --- a/model_compression_toolkit/qat/keras/README.md +++ b/model_compression_toolkit/qat/keras/README.md @@ -1,38 +1,41 @@ # QAT Quantizers ## Introduction -[`BaseKerasQATTrainableQuantizer`](./quantizer/base_keras_qat_quantizer.py) is an interface that utilizes the Quantization Infrastructure's [`BaseKerasTrainableQuantizer`](../../trainable_infrastructure/keras/base_keras_quantizer.py) class to enable easy development of quantizers dedicated to Quantization-Aware Training (QAT). -All available training types for QAT are defined in the Enum `TrainingMethod`. See [Training Methods for QAT](./quantizer/README.md) for more details. +All available training types for QAT are defined in the Enum [`TrainingMethod`](../../trainable_infrastructure/common/training_method.py). +A trainable quantizer can be Weights Quantizer or Activation Quantizer. +Any Activation Quantizer defined in [Activation Quantizers](../../trainable_infrastructure/keras/activation_quantizers) can be used for QAT. ## Make your own Keras trainable quantizers -A trainable quantizer can be Weights Quantizer or Activation Quantizer. In order to make your new quantizer you need to create your quantizer class, `MyTrainingQuantizer` and do as follows: - - `MyTrainingQuantizer` should inherit from `BaseKerasQATTrainableQuantizer`. - - `MyTrainingQuantizer` should have [`init`](../../trainable_infrastructure/common/base_trainable_quantizer.py) function that gets `quantization_config` which is [`NodeWeightsQuantizationConfig`](../../core/common/quantization/node_quantization_config.py#L228) if you choose to implement weights quantizer or [`NodeActivationQuantizationConfig`](../../core/common/quantization/node_quantization_config.py#L63) if you choose activation quantizer. - - Implement [`initialize_quantization`](../../trainable_infrastructure/common/base_trainable_quantizer.py) where you can define your parameters for the quantizer. - - Implement [`__call__`](../../trainable_infrastructure/common/base_trainable_quantizer.py) method to quantize the given inputs while training. This is your custom quantization itself. - - Implement [`convert2inferable`](../../trainable_infrastructure/common/base_trainable_quantizer.py) method. This method exports your quantizer for inference (deployment). For doing that you need to choose one of the available Inferable Quantizers from the [MCT Quantizers](https://github.com/sony/mct_quantizers) package, according to the target when implementing `convert2inferable`, and set your learned quantization parameters there. - - Decorate `MyTrainingQuantizer` class with the `@mark_quantizer` decorator (provided by the [MCT Quantizers](https://github.com/sony/mct_quantizers) package) and choose the appropriate properties to set for your quantizer. The "identifier" property for the decorator should be of type `TrainingMethod` enum. See explanation about `@mark_quantizer` and how to use it under the [Pytorch Quantization Infrastructure](../../trainable_infrastructure/keras/README.md). + - `MyTrainingQuantizer` should inherit from `BaseKerasQATWeightTrainableQuantizer` for weights quantizer or `BaseKerasActivationTrainableQuantizer` for activation quantizer + - `MyTrainingQuantizer` should have `__init__` method that accepts `quantization_config` of type `TrainableQuantizerWeightsConfig` for weights quantizer or `TrainableQuantizerActivationConfig` for activation quantizer. + - Implement `initialize_quantization` where you can define your parameters for the quantizer. + - Implement `__call__` method to quantize the given inputs while training. This is your custom quantization itself. + - Implement `convert2inferable` method. This method exports your quantizer for inference (deployment). For doing that you need to choose one of the available Inferable Quantizers from the [MCT Quantizers](https://github.com/sony/mct_quantizers) package, according to the target when implementing `convert2inferable`, and set your learned quantization parameters there. + - Decorate `MyTrainingQuantizer` class with the `@mark_quantizer` decorator (provided by the [MCT Quantizers](https://github.com/sony/mct_quantizers) package) and choose the appropriate properties to set for your quantizer. The "identifier" property for the decorator should be of the type `TrainingMethod` enum. See explanation about `@mark_quantizer` and how to use it under the [Keras Quantization Infrastructure](../../trainable_infrastructure/keras/README.md). + + ## Example: Symmetric Weights Quantizer To create custom `MyWeightsTrainingQuantizer` which is a symmetric weights training quantizer you need to set -`qi.QuantizationTarget.Weights` as target and `qi.QuantizationMethod.SYMMETRIC` as method. +`QuantizationTarget.Weights` as target and `QuantizationMethod.SYMMETRIC` as method. Assume that the quantizer has a new training method called `MyTrainig` which is defined in the `TrainingMethod` Enum. ```python import tensorflow as tf -from model_compression_toolkit.trainable_infrastructure import TrainingMethod + +from model_compression_toolkit.trainable_infrastructure import TrainingMethod, TrainableQuantizerWeightsConfig from mct_quantizers import mark_quantizer, QuantizationMethod, QuantizationTarget -from model_compression_toolkit.qat.keras.quantizer.base_keras_qat_quantizer import BaseKerasQATTrainableQuantizer +from mct_quantizers.keras.quantizers import WeightsUniformInferableQuantizer +from model_compression_toolkit.qat.keras.quantizer.base_keras_qat_weight_quantizer import BaseKerasQATWeightTrainableQuantizer NEW_PARAM = "new_param_name" - @mark_quantizer(quantization_target=QuantizationTarget.Weights, quantization_method=[QuantizationMethod.SYMMETRIC], identifier=TrainingMethod.MyTraining) -class MyWeightsTrainingQuantizer(BaseKerasQATTrainableQuantizer): - def __init__(self, quantization_config: NodeWeightsQuantizationConfig): +class MyWeightsTrainingQuantizer(BaseKerasQATWeightTrainableQuantizer): + def __init__(self, quantization_config: TrainableQuantizerWeightsConfig): super(MyWeightsTrainingQuantizer, self).__init__(quantization_config) # Define your new params here: self.new_param = ... @@ -56,28 +59,30 @@ class MyWeightsTrainingQuantizer(BaseKerasQATTrainableQuantizer): return quantized_inputs def convert2inferable(self): - return qi.WeightsUniformInferableQuantizer(...) + return WeightsUniformInferableQuantizer(...) + ``` ## Example: Symmetric Activations Quantizer -To create custom `MyActivationsTrainingQuantizer` which is a symmetric activations training quantizer you need to set `qi.QuantizationTarget.Activation` as target and `qi.QuantizationMethod.SYMMETRIC` as method. +To create custom `MyActivationsTrainingQuantizer` which is a symmetric activations training quantizer you need to set `QuantizationTarget.Activation` as target and `QuantizationMethod.SYMMETRIC` as method. Assume that the quantizer has a new training method called `MyTrainig` which is defined in the `TrainingMethod` Enum. ```python import tensorflow as tf -NEW_PARAM = "new_param_name" -from model_compression_toolkit.trainable_infrastructure import TrainingMethod +from mct_quantizers.keras.quantizers import ActivationUniformInferableQuantizer +from model_compression_toolkit.trainable_infrastructure import TrainingMethod, TrainableQuantizerActivationConfig from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod from mct_quantizers import mark_quantizer, QuantizationTarget -from model_compression_toolkit.qat.keras.quantizer.base_keras_qat_quantizer import BaseKerasQATTrainableQuantizer +from model_compression_toolkit.qat.keras.quantizer.base_keras_qat_weight_quantizer import BaseKerasQATWeightTrainableQuantizer +NEW_PARAM = "new_param_name" @mark_quantizer(quantization_target=QuantizationTarget.Activation, quantization_method=[QuantizationMethod.SYMMETRIC], identifier=TrainingMethod.TrainingMethod) -class MyActivationsTrainingQuantizer(BaseKerasQATTrainableQuantizer): - def __init__(self, quantization_config: NodeActivationQuantizationConfig): +class MyActivationsTrainingQuantizer(BaseKerasQATWeightTrainableQuantizer): + def __init__(self, quantization_config: TrainableQuantizerActivationConfig): super(MyActivationsTrainingQuantizer, self).__init__(quantization_config) # Define your new params here: self.new_param = ... @@ -101,5 +106,5 @@ class MyActivationsTrainingQuantizer(BaseKerasQATTrainableQuantizer): return quantized_inputs def convert2inferable(self): - return qi.ActivationUniformInferableQuantizer(...) + return ActivationUniformInferableQuantizer(...) ```