Add A/B test against V1 range learning quantizer

Signed-off-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>
quic · Feb 2, 2024 · 1db6857 · 1db6857
1 parent ae77cf7
commit 1db6857
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/quantsim.py b/TrainingExtensions/torch/src/python/aimet_torch/quantsim.py
@@ -765,7 +765,12 @@ def replace_wrappers_for_quantize_dequantize(self):
         """
         if self._quant_scheme == QuantScheme.training_range_learning_with_tf_init or self._quant_scheme == \
                 QuantScheme.training_range_learning_with_tf_enhanced_init:
-            device = utils.get_device(self.model)
+            try:
+                device = utils.get_device(self.model)
+            except StopIteration:
+                # Model doesn't have any parameter.
+                # Set device to cpu by default.
+                device = torch.device('cpu')
 
             self._replace_quantization_wrapper(self.model, device)
 

diff --git a/TrainingExtensions/torch/test/python/experimental/v2/ab_test/test_quantsim_logits.py b/TrainingExtensions/torch/test/python/experimental/v2/ab_test/test_quantsim_logits.py
@@ -198,6 +198,7 @@ def set_seed(seed):
 
 
 @pytest.mark.parametrize('quant_scheme', [QuantScheme.post_training_tf,
+                                          QuantScheme.training_range_learning_with_tf_init,
                                           # QuantScheme.post_training_percentile, # TODO: not implemented
                                           # QuantScheme.training_range_learning_with_tf_init, # TODO: not implemented
                                          ])
@@ -288,4 +289,3 @@ def test_multi_output(self, quant_scheme, seed):
         model = models_to_test.ModelWith5Output()
         dummy_input = torch.randn(1, 3, 224, 224)
         self.check_qsim_logit_consistency(CONFIG_DEFAULT, quant_scheme, model, dummy_input)
-