Quantization-Aware-Training with Sparsity and Clustering Preservation

[LesBellArm]

System information

• TensorFlow version (you are using): TF.2.3.0
• Are you willing to contribute it (Yes/No): Yes

Motivation

Sparse (pruned), clustered and quantized models can help to reduce the compressed model sizes for deployment to devices that may be constrained by memory size and processing power. These model optimizations are already supported by the Model Optimization Toolkit and their benefits are already documented elsewhere. Combining these optimizations in a single model can, currently, only be done at the post-training stage, with a subsequent loss of accuracy. Applying fine-tuning to the resulting model, in order to recover the accuracy, will modify the weights so that the sparsity and clustering benefits are lost.

The use case is to be able to perform Quantization-Aware-Training without losing the sparsity or clustering properties of models that have already been optimized by these other techniques, while benefiting from the improved accuracy of the fine-tune training process.

Describe the feature

This feature proposes the addition of new Quantizer classes (as used in Quantization-Aware-Training) along with their associated QuantizeConfig classes. Specifically, the following new Quantizer derived classes should be added:

• Prune Quantizer - To preserve zero weights during the training process.
• Cluster Quantizer - To preserve the number of unique weights (centroids) during training.
• Prune and Cluster Quantizer - To preserve the zero weights and the number of unique weights during training.
• QuantizeConfig derived classes for each of the above.

Describe how the feature helps achieve the use case

Quantizer classes that preserve the sparsity and clustering properties of a model that has already been optimized by these other techniques will allow the fine-tuning of the model during the quantization process, which will help to maintain the accuracy of the optimized model.

Describe how existing APIs don't satisfy your use case (optional if obvious)

It is not currently possible to preserve the sparsity or clustering of a model during the training process.

[alanchiao]

@nutsiepully, @liyunlu0618

[alanchiao]

Discussed this on the call but to make note:

1. It'd be necessary to demonstrate that there are a significant number of cases where cascading all three can lead to better results (whether size or latency) for the same given accuracy. To start, it may be worth just trying to combine two - since then that is a step forward already in being able to combine quantization-aware training with another technique, not just post-training quantization. This can then be weighed against the cost of maintenance.

A convincing thing would be take some of the latest models (from TFHub and TFLite) and produce models that exceed them significantly. I will check with others to see what the best test models are in peoples' minds.

2. It's nice regardless that this is a modular approach that doesn't require modifying any of the individual techniques' code. Given that this is taking the cascaded approach as opposed to the combining N techniques at once, this seems to be the most sensible thing I can think of today, especially to get started with running experiments to determine the value.

Post-meeting:

1. realized that that even if you create QuantizeConfig/Quantizers so that the exported inference graph would be exactly the same as when using just Quantization Aware Training, there are limitations today. In particular, the handling of cross-layer aspects (layer + BN and layer1 + layer 2 ->Concat) wouldn't happen. You can try having a layer1 with the same QuantizeConfig used within our default behavior, and it will throw a "layer 1 unsupported" error.
   The relevant code is in the transforms, such as here.

[psunn]

@alanchiao
Hi alan, I'm trying to enable Prune preserve QAT in MOT. I need some suggestions regarding QAT API.
Current QAT API:
per layer:
tfmot.quantization.keras.quantize_annotate_layer(keras_layer)
per model:
tfmot.quantization.keras.quantize_model(keras_model)

For Prune-preserve QAT:
would it make sense to add boolean param:
per layer:
tfmot.quantization.keras.quantize_annotate_layer(keras_layer, prune_preserve=True)
per model:
tfmot.quantization.keras.quantize_model(keras_model, prune_preserve=True)
(I lean towards this more, as I consider prune_preserve as a selective mode for QAT)

or perhaps (based on #520 conversation) :
per layer:
tfmot.quantization.keras.experimental_prune_preserve.quantize_annotate_layer(keras_layer)
per model:
tfmot.quantization.keras.experimental_prune_preserve.quantize_model(keras_model)

or others?

[alanchiao]

@nutsiepully, @liyunlu0618 for visibility

[nutsiepully]

@psunn - Thanks for the suggestions. These are all fundamental changes to the API, and wouldn't quite work.

We designed the QAT API to be flexible enough to handle changes. We haven't exposed some of this yet, but we could to make your work easier. The plan was for quantize_model and quantize_apply to expose a QuantizationScheme param which allows users to customize both the transforms, and the registry.

You could easily use that to alter behavior how you want.

[psunn]

@wwwind for visibility

[psunn]

@nutsiepully Thanks for sharing the ongoing work and future changes.
I guess we can have a flag QuantizationScheme.preserve_sparsity, when QAT API updated.
To enable prune-preserve QAT, user can:

QuantizationScheme.preserve_sparsity = True
tfmot.quantization.keras.quantize_model(keras_model, QuantizationScheme)

For experimental and feature approval, probably this is ok for now:
tfmot.quantization.keras.experimental_prune_preserve.quantize_annotate_layer(keras_layer)

what do you think?

[nutsiepully]

Hi @psunn,

The goal of the API is to allow users to plug in their own QuantizationScheme objects. Within them you can add whatever field you want, and you can also modify behavior how you want.

It'll simply be tfmot.quantization.keras.quantize_model(keras_model, custom_quantization_scheme). The custom_quantization_scheme is where you can choose to do what you want.

[psunn]

Hi @nutsiepully, thanks for confirm.

Do you know when custom_quantization_scheme will be available with master branch, is it available now with a public branch so that I can take a look?

[nutsiepully]

@psunn - I haven't added it yet since the need didn't come up. But I can add it if you need it.

[psunn]

#568 raised for QAT with sparsity preservation.