[quant] Remove get_qparams in Observers (#38435)

jerryzh168 · facebook-github-bot · commit 86397f6b2459 · 2020-05-18T20:49:33.000-07:00
Summary: Pull Request resolved: #38435 Test Plan: Imported from OSS Differential Revision: D21597835 Pulled By: jerryzh168 fbshipit-source-id: 88a8dd110db5586509bf98fa6712290f1756c272
diff --git a/torch/csrc/jit/passes/quantization/insert_quant_dequant.cpp b/torch/csrc/jit/passes/quantization/insert_quant_dequant.cpp
@@ -339,17 +339,17 @@ void RemoveRedundantQuantizationOps(std::shared_ptr<Graph>& graph) {
   rewriter.runOnGraph(graph, filter);
 }
 
-void checkGetQParamsResult(const IValue& qparams) {
+void checkCalculateQParamsResult(const IValue& qparams) {
   TORCH_CHECK(
       qparams.isTuple(),
-      "`get_qparams` function is expected to return a "
+      "`calculate_qparams` function is expected to return a "
       "Tuple, but got:",
       qparams.tagKind());
   auto tp = qparams.toTuple();
   TORCH_CHECK(
-      tp->elements().size() == 2 || tp->elements().size() == 3,
-      "`get_qparams` function is expected to return a "
-      "Tuple of size 2 or 3, got Tuple of size ",
+      tp->elements().size() == 2,
+      "`calculate_qparams` function is expected to return a "
+      "Tuple of size 2, got Tuple of size ",
       tp->elements().size());
   // Expect first two elements of the tuple to be Tensor
   for (size_t i = 0; i < 2; ++i) {
@@ -360,15 +360,6 @@ void checkGetQParamsResult(const IValue& qparams) {
         " has type: ",
         tp->elements()[i].tagKind());
   }
-  // Expect the third elements of the tuple to be int
-  if (tp->elements().size() == 3) {
-    TORCH_CHECK(
-        tp->elements()[2].isInt(),
-        "Element of Tuple is expected to be int, but element ",
-        2,
-        " has type: ",
-        tp->elements()[2].tagKind());
-  }
 }
 
 class InsertQuantDeQuantHelper {
@@ -567,9 +558,9 @@ std::tuple<c10::QScheme, QParamVector> InsertQuantDeQuantHelper::
       v->debugName(),
       " exists.");
   auto observer_module = module.attr(observer_name.value()).toModule();
-  auto get_qparams = observer_module.get_method("get_qparams");
-  IValue result = get_qparams(std::vector<IValue>());
-  checkGetQParamsResult(result);
+  auto calculate_qparams = observer_module.get_method("calculate_qparams");
+  IValue result = calculate_qparams(std::vector<IValue>());
+  checkCalculateQParamsResult(result);
   auto scalar_type = observer_module.attr("dtype");
   TORCH_CHECK(
       scalar_type.toScalarType() != at::ScalarType::Undefined,
@@ -582,9 +573,10 @@ std::tuple<c10::QScheme, QParamVector> InsertQuantDeQuantHelper::
   QParamVector qparams;
   auto qscheme = observer_module.attr("qscheme").toQScheme();
   if (isPerChannel(qscheme)) {
+    auto axis = observer_module.attr("ch_axis");
     qparams.push_back(std::make_pair("_scale", scale));
     qparams.push_back(std::make_pair("_zero_point", zero_point));
-    qparams.push_back(std::make_pair("_axis", tp->elements()[2].toInt()));
+    qparams.push_back(std::make_pair("_axis", axis.toInt()));
   } else {
     qparams.push_back(std::make_pair("_scale", scale.item<double>()));
     qparams.push_back(
diff --git a/torch/quantization/observer.py b/torch/quantization/observer.py
@@ -66,14 +66,6 @@ def forward(self, x):
     def calculate_qparams(self, **kwargs):
         pass
 
-    # Returns all quantization parameters that's needed
-    # for a quantize function call
-    # For instance, per channel obsserver will return
-    # scales, zero_points and axis
-    @abstractmethod
-    def get_qparams(self, **kwargs):
-        pass
-
     with_args = classmethod(_with_args)
 
 
@@ -194,10 +186,6 @@ def _calculate_qparams(self, min_val, max_val):
 
         return scale, zero_point
 
-    @torch.jit.export
-    def get_qparams(self):
-        r"""Get all quantization parameters needed for quantize call"""
-        return self.calculate_qparams()
 
 class MinMaxObserver(_ObserverBase):
     r"""Observer module for computing the quantization parameters based on the
@@ -546,11 +534,6 @@ def _forward(self, x_orig):
     def calculate_qparams(self):
         return self._calculate_qparams(self.min_vals, self.max_vals)
 
-    @torch.jit.export
-    def get_qparams(self):
-        scales, zero_points = self.calculate_qparams()
-        return scales, zero_points, self.ch_axis
-
     def extra_repr(self):
         return "min_val={}, max_val={}".format(self.min_vals, self.max_vals)
 
@@ -966,12 +949,10 @@ def __init__(self, dtype=torch.float16):
     def forward(self, x):
         return x
 
+    @torch.jit.export
     def calculate_qparams(self):
         raise Exception("calculate_qparams should not be called for NoopObserver")
 
-    def get_qparams(self):
-        return self.calculate_qparams()
-
 
 # Restrict activations to be in the range (0,127)
 default_observer = MinMaxObserver.with_args(reduce_range=True)
