Changes for TF quant analyzer (#2610)

Signed-off-by: Chetan Gulecha <quic_cgulecha@quicinc.com>
quic · Dec 19, 2023 · e8a8ff7 · e8a8ff7
1 parent f3e1090
commit e8a8ff7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/Docs/api_docs/tensorflow_quantization.rst b/Docs/api_docs/tensorflow_quantization.rst
@@ -4,6 +4,7 @@ AIMET TensorFlow Quantization APIs
 
 AIMET Quantization for TensorFlow provides the following functionality
    - :ref:`Quantization Simulation<api-tf-quantsim>`: Allows ability to simulate inference and training on quantized hardware
+   - :ref:`QuantAnalyzer<api-tensorflow-quant-analyzer>`: Analyzes the model and points out sensitive ops to quantization
    - :ref:`Adaptive Rounding<api-tf-adaround>`: Post-training quantization technique to optimize rounding of weight tensors
    - :ref:`Cross-Layer Equalization<api-tf-cle>`: Post-training quantization technique to equalize layer parameters
    - :ref:`Bias Correction<api-tf-bias-correction>`: Post-training quantization technique to correct shift in layer outputs due to quantization noise

diff --git a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/quant_analyzer.py b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/quant_analyzer.py
@@ -402,7 +402,7 @@ def _compute_mse_loss(self, sim: QuantizationSimModel, output_op_name) -> float:
 
             # Collect output activation data from quant sim op
             feed_dict = create_input_feed_dict(sim.session.graph, self._start_op_names, model_inputs)
-            quant_op = sim.session.graph.get_operation_by_name(output_op_name)
+            quant_op = sim.session.graph.get_operation_by_name(output_op_name + "_quantized")
             quantized_out_data = sim.session.run(quant_op.outputs[0], feed_dict=feed_dict)
 
             # Calculate MSE loss