TFLite conversion with quantization after quantization-aware training of simple CNN model fails

[kalaluthien]

Describe the bug

Simple MNIST cnn model inspired by tfmot functionality test and examples, failed to convert with post-training integer quantization.

System information

TensorFlow installed from (source or binary): binary

• pip install tf-nightly

TensorFlow version: 2.2.0 (latest)

TensorFlow Model Optimization version: 0.3.0 (latest)

Python version: 3.6 (colab) & 3.7 (local)

Describe the expected behavior
QAT + PTQ on simple model works well.

Describe the current behavior
Failed to quantize model at convert phase...

For local machine (Python3.7):

Traceback (most recent call last):
  File "/home/mumu/.miniconda3/envs/tf-nightly-2.2.0rc2/bin/toco_from_protos", line 8, in <module>
    sys.exit(main())
  File "/home/mumu/.miniconda3/envs/tf-nightly-2.2.0rc2/lib/python3.7/site-packages/tensorflow/lite/toco/python/toco_from_protos.py", line 93, in main
    app.run(main=execute, argv=[sys.argv[0]] + unparsed)
  File "/home/mumu/.miniconda3/envs/tf-nightly-2.2.0rc2/lib/python3.7/site-packages/tensorflow/python/platform/app.py", line 40, in run
    _run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
  File "/home/mumu/.miniconda3/envs/tf-nightly-2.2.0rc2/lib/python3.7/site-packages/absl/app.py", line 299, in run
    _run_main(main, args)
  File "/home/mumu/.miniconda3/envs/tf-nightly-2.2.0rc2/lib/python3.7/site-packages/absl/app.py", line 250, in _run_main
    sys.exit(main(argv))
  File "/home/mumu/.miniconda3/envs/tf-nightly-2.2.0rc2/lib/python3.7/site-packages/tensorflow/lite/toco/python/toco_from_protos.py", line 56, in execute
    enable_mlir_converter)
Exception: /home/mumu/.miniconda3/envs/tf-nightly-2.2.0rc2/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py:744:34: error: 'tf.FakeQuantWithMinMaxVarsPerChannel' op operand #0 must be tensor of 32-bit float values, but got 'tensor<5x5x1x32x!quant.uniform<i8:f32, 6.819487202401255E-4:-3>>'

For gist (Python3.6):

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-6-7ddb90644e60> in <module>()
     14 
     15     with quantize.quantize_scope():
---> 16         tflite_model = converter.convert()
     17         converted_models.append((tflite_model, model.name))

/usr/local/lib/python3.6/dist-packages/tensorflow/lite/python/lite.py in convert(self)
    683     if quant_mode.post_training_int8_no_float():
    684       result = self._calibrate_quantize_model(result, constants.FLOAT,
--> 685                                               constants.FLOAT, False)
    686     elif quant_mode.post_training_int8_allow_float():
    687       result = self._calibrate_quantize_model(result, constants.FLOAT,

/usr/local/lib/python3.6/dist-packages/tensorflow/lite/python/lite.py in _calibrate_quantize_model(self, result, inference_input_type, inference_output_type, allow_float)
    335       return calibrate_quantize.calibrate_and_quantize(
    336           self.representative_dataset.input_gen, inference_input_type,
--> 337           inference_output_type, allow_float)
    338 
    339   def _is_unknown_shapes_allowed(self, fp32_execution):

/usr/local/lib/python3.6/dist-packages/tensorflow/lite/python/optimize/calibrator.py in calibrate_and_quantize(self, dataset_gen, input_type, output_type, allow_float, resize_input)
     91     return self._calibrator.QuantizeModel(
     92         np.dtype(input_type.as_numpy_dtype()).num,
---> 93         np.dtype(output_type.as_numpy_dtype()).num, allow_float)
     94 
     95   def calibrate_and_quantize_single(self,

RuntimeError: Max and min for dynamic tensors should be recorded during calibration

Code to reproduce the issue
See gist

Screenshots
None

Additional context
None

[alanchiao]

For combining QAT with post-training full-integer quantization, you may see code examples, but that has not launched yet and is not ready for usage. Once it's launched and there is official documentation, you can start filing issues and bugs against it.

[grayhong]

@alanchiao Any plan on this?