Skip to content
Branch: master
Find file Copy path
Find file Copy path
9 contributors

Users who have contributed to this file

@suharshs @tensorflower-gardener @Randl @siju-samuel @raghuraman-k @unnir @jianlijianli @lamberta @angersson
165 lines (130 sloc) 7.14 KB

Quantization-aware training

Quantization-aware model training ensures that the forward pass matches precision for both training and inference. There are two aspects to this:

  • Operator fusion at inference time are accurately modeled at training time.
  • Quantization effects at inference are modeled at training time.

For efficient inference, TensorFlow combines batch normalization with the preceding convolutional and fully-connected layers prior to quantization by folding batch norm layers{:.external}.

The quantization error is modeled using fake quantization nodes to simulate the effect of quantization in the forward and backward passes. The forward-pass models quantization, while the backward-pass models quantization as a straight-through estimator. Both the forward- and backward-pass simulate the quantization of weights and activations. Note that during back propagation, the parameters are updated at high precision as this is needed to ensure sufficient precision in accumulating tiny adjustments to the parameters.

Additionally, the minimum and maximum values for activations are determined during training. This allows a model trained with quantization in the loop to be converted to a fixed point inference model with little effort, eliminating the need for a separate calibration step.

Since it's difficult to add these fake quantization operations to all the required locations in the model, there's a function available that rewrites the training graph. To create a fake quantized training graph:

# Build forward pass of model.
loss = tf.losses.get_total_loss()

# Call the training rewrite which rewrites the graph in-place with
# FakeQuantization nodes and folds batchnorm for training. It is
# often needed to fine tune a floating point model for quantization
# with this training tool. When training from scratch, quant_delay
# can be used to activate quantization after training to converge
# with the float graph, effectively fine-tuning the model.
g = tf.get_default_graph()

# Call backward pass optimizer as usual.
optimizer = tf.train.GradientDescentOptimizer(learning_rate)

The rewritten eval graph is non-trivially different from the training graph since the quantization ops affect the batch normalization step. Because of this, we've added a separate rewrite for the eval graph:

# Build eval model
logits = tf.nn.softmax_cross_entropy_with_logits_v2(...)

# Call the eval rewrite which rewrites the graph in-place with
# FakeQuantization nodes and fold batchnorm for eval.
g = tf.get_default_graph()

# Save the checkpoint and eval graph proto to disk for freezing
# and providing to TFLite.
with open(eval_graph_file, ‘w’) as f:
saver = tf.train.Saver(), checkpoint_name)

Methods to rewrite the training and eval graphs are an active area of research and experimentation. Although rewrites and quantized training might not work or improve performance for all models, we are working to generalize these techniques.

Generating fully-quantized models

The previously demonstrated after-rewrite eval graph only simulates quantization. To generate real fixed-point computations from a trained quantization model, convert it to a fixed-point kernel. TensorFlow Lite supports this conversion from the graph resulting from create_eval_graph.

First, create a frozen graph that will be the input for the TensorFlow Lite toolchain:

freeze_graph \
  --input_graph=eval_graph_def.pb \
  --input_checkpoint=checkpoint \
  --output_graph=frozen_eval_graph.pb --output_node_names=outputs

Provide this to the TensorFlow Lite Optimizing Converter (TOCO) to get a fully-quantized TensorFlow Lite model:

toco \
  --input_file=frozen_eval_graph.pb \
  --output_file=tflite_model.tflite \
  --input_format=TENSORFLOW_GRAPHDEF --output_format=TFLITE \
  --inference_type=QUANTIZED_UINT8 \
  --input_shape="1,224, 224,3" \
  --input_array=input \
  --output_array=outputs \
  --std_value=127.5 --mean_value=127.5

See the documentation for tf.contrib.quantize and TensorFlow Lite.

Quantized accuracy results

The following are results of training some popular CNN models (Mobilenet-v1, Mobilenet-v2, and Inception-v3) using this tool:

Model Top-1 Accuracy:
Floating point
Top-1 Accuracy:
Fixed point: 8 bit weights and activations
Table 1: Top-1 accuracy of floating point and fully quantized CNNs on Imagenet Validation dataset.

Our pre-trained models are available in the TensorFlow Lite model repository. The code used to generate these models is available.

These rewrites are an active area of research and experimentation, so the rewrites and quantized training will likely not work across all models, though we hope to work towards generalizing these techniques.

[1] B.Jacob et al., "Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference",


[3] Y.Bengio et al., "Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation",

You can’t perform that action at this time.