Fix TF QAT UT issues (#266)

Signed-off-by: zehao-intel <[email protected]>

Author: zehao-intel

neural_compressor/adaptor/tensorflow.py

@@ -525,6 +525,17 @@ def quantize(self, tune_cfg, model, data_loader, q_func=None):
         Returns:
             tf.compat.v1.GraphDef: the quantized model
         """
+        if self.approach == "quant_aware_training":
+            assert q_func is not None, "quantization aware training mode \
+                is not configured correctly"
+
+            from neural_compressor.experimental import common
+            qat_model = q_func(model)
+
+            return self.convert(common.Model(qat_model), 'QAT', 'default')
+
+        assert q_func is None, \
+            "post-training quantization mode is not support calibration function for Tensorflow!"
         self.tuning_cfg_to_fw(tune_cfg)
         logger.debug("Dump quantization configurations:")
         logger.debug(self.quantize_config)

test/quantization/test_tensorflow_qat.py

@@ -88,6 +88,7 @@ def test_qat(self):
         compression_manager = training.prepare_compression('./baseline_model', config)
         compression_manager.callbacks.on_train_begin()
 
+        q_aware_model = compression_manager.model
         # `quantize_model` requires a recompile.
         q_aware_model.compile(optimizer='adam',
                                 loss=tf.keras.losses.SparseCategoricalCrossentropy(