Temporarily use fp32 to recompute encodings

Signed-off-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>

TrainingExtensions/torch/src/python/aimet_torch/tensor_quantizer.py

@@ -698,14 +698,13 @@ def encoding(self) -> Union[None, libpymo.TfEncoding, List[libpymo.TfEncoding]]:
 
         :return: encoding(s).
         """
-        # pylint:disable=protected-access
-        if self.enabled:
-            if self.bitwidth == 32 or self.data_type == QuantizationDataType.float:
-                return None
-            encoding = self._compute_updated_encoding()
-            return encoding
+        if not self.enabled:
+            return None
 
-        return None
+        if self.bitwidth == 32 or self.data_type == QuantizationDataType.float:
+            return None
+
+        return self._compute_updated_encoding()
 
     @encoding.setter
     def encoding(self, encoding: Union[libpymo.TfEncoding, List[libpymo.TfEncoding]]):
@@ -809,25 +808,32 @@ def _compute_updated_encoding(self) -> Union[libpymo.TfEncoding, List[libpymo.Tf
         encoding_min = getattr(self.wrapper_ref, self.name + '_encoding_min')
         encoding_max = getattr(self.wrapper_ref, self.name + '_encoding_max')
 
-        encodings = []
-        for minimum, maximum in zip(encoding_min, encoding_max):
-            tf_encoding = libpymo.TfEncoding()
-            scale, offset = self.compute_scaling_offset(minimum, maximum)
+        if encoding_min is None or encoding_max is None:
+            return None # Encoding is not inititialized yet
+
+        scale, offset = self.compute_scaling_offset(encoding_min.float(), encoding_max.float())
+        assert scale is not None
+        assert offset is not None
 
+        if not self.use_symmetric_encodings or self.is_unsigned_symmetric:
             # Calculate 'min' and 'max' based on 'delta' and 'offset'
             #   because offset was adjusted so that zero must be quantizable
             #   in the case of asymmetric or unsigned symmetric quantization
             #   Please refer quantization_utils.cpp if you need
-            if not self.use_symmetric_encodings or self.is_unsigned_symmetric:
-                tf_encoding.min = scale * offset
-                # NOTE: We want to calculate: max = delta * numSteps + min
-                #   To avoid numerical accuracy issues on Linaro, we simplify the math.
-                tf_encoding.max = maximum - minimum + tf_encoding.min
-            else:
-                tf_encoding.min = minimum
-                tf_encoding.max = maximum
+            # NOTE: We want to calculate: max = delta * numSteps + min
+            #   To avoid numerical accuracy issues on Linaro, we simplify the math.
+            adjusted_min = scale * offset
+            encoding_max = encoding_max - encoding_min + adjusted_min
+            encoding_min = adjusted_min
 
-            tf_encoding.offset, tf_encoding.delta, tf_encoding.bw = offset, scale, self.bitwidth
+        encodings = []
+        for min_, max_, scale_, offset_ in zip(encoding_min, encoding_max, scale, offset):
+            tf_encoding = libpymo.TfEncoding()
+            tf_encoding.min = min_
+            tf_encoding.max = max_
+            tf_encoding.delta = scale_
+            tf_encoding.offset = offset_
+            tf_encoding.bw = self.bitwidth
             encodings.append(tf_encoding)
 
         # TODO: Remove when using only sequence of encodings (Done for backward compatibility)

TrainingExtensions/torch/test/python/test_tensor_quantizer.py

@@ -36,6 +36,7 @@
 # =============================================================================
 
 import random
+import numpy as np
 import pytest
 import torch
 import aimet_common.libpymo as libpymo
@@ -394,3 +395,23 @@ def test_non_contiguous_input_tensor(self, quantizer):
         out_strided = quantizer.quantize_dequantize(tensor[..., ::2, :],
                                                     quantizer.round_mode)
         assert torch.allclose(out_contiguous, out_strided)
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
+    def test_learned_grid_encoding_getter(self, dtype):
+        conv = torch.nn.Conv2d(1, 1, 1)
+        quant_wrapper = LearnedGridQuantWrapper(conv, round_mode=libpymo.RoundingMode.ROUND_NEAREST,
+                                                quant_scheme=QuantScheme.training_range_learning_with_tf_init,
+                                                is_output_quantized=True, activation_bw=16,
+                                                weight_bw=8, device="cuda:0")
+        enc = libpymo.TfEncoding()
+        enc.bw, enc.max, enc.min = 16, 0.4, -0.98
+
+        quant_wrapper.output_quantizers[0].enabled = True
+        quant_wrapper.output_quantizers[0].encoding = enc
+
+        e = quant_wrapper.to(dtype=dtype).output_quantizers[0].encoding
+        rtol=1e-3
+        assert np.isclose(e.max, 0.4,   rtol=rtol)
+        assert np.isclose(e.min, -0.98, rtol=rtol)
+        assert np.isclose(e.delta, (e.max-e.min)/(2**e.bw-1), rtol=rtol)
+        assert np.isclose(e.offset, e.min/e.delta, rtol=rtol)