[quant][graphmode] use whitelist for selecting observed values

jerryzh168 · jerryzh168 · commit 34a1c47423a9 · 2019-09-11T17:36:32.000-07:00
Summary: Previously we observe all the Tensor values, but what we want is actually observing only the ones that can be quantized. Test Plan: python test/test_jit.py python test/test_quantizer.py Reviewers: pt1quant Subscribers: Tasks: Tags: ghstack-source-id: 68e19e6 Pull Request resolved: #25974
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -876,7 +876,6 @@ def test_recursive_cse(self):
         FileCheck().run(input_str, graph)
 
     @_tmp_donotuse_dont_inline_everything
-    @unittest.skip("temoprarily disable the test case, it will pass in later PR")
     def test_insert_observers(self):
         class Observer(torch.nn.Module):
             def __init__(self):
@@ -933,7 +932,6 @@ def get_forward_graph(m):
                    .run(str(m._c._get_module("conv")._get_method('conv2d_forward').graph))
 
     @_tmp_donotuse_dont_inline_everything
-    @unittest.skip("temoprarily disable the test case, it will pass in later PR")
     def test_insert_observers_child_qconfig(self):
         class Observer(torch.nn.Module):
             def __init__(self):
@@ -1004,8 +1002,6 @@ def get_forward(c):
         check_observed(get_forward(m._c._get_module('sub')._get_module('linear')).graph)
 
     @_tmp_donotuse_dont_inline_everything
-    @unittest.skip("temoprarily disable the test since \
-    I want to put the insert_quant_dequant changes in a separate PR")
     def test_insert_observers_skip_values(self):
         import torch.nn.functional as F
 
@@ -1043,7 +1039,6 @@ def get_forward(m):
         def test_module(module, relu_call, num_observers):
             m = torch.jit.script(module())
             observer = torch.jit.script(Observer())
-
             torch._C._jit_pass_constant_propagation(get_forward(m).graph)
             torch._C._jit_pass_constant_propagation(m._c._get_module('conv')._get_method('conv2d_forward').graph)
             qconfig_dict = {
diff --git a/torch/csrc/jit/passes/quantization.cpp b/torch/csrc/jit/passes/quantization.cpp
@@ -64,8 +64,40 @@ graph(%self, %input):
   }
 }
 
-static bool outputsNeedToBeObserved(Node* n) {
-  return n->kind() != prim::Constant;
+bool nodeQuantizable(Node* n) {
+  static std::vector<Symbol> aten_funcs = {
+    Symbol::aten("conv2d"),
+    Symbol::aten("linear")
+  };
+  bool is_quantizable = std::find(aten_funcs.begin(), aten_funcs.end(), n->kind()) != aten_funcs.end();
+  static std::vector<std::string> call_funcs = {
+    "linear",
+    "relu"
+  };
+  if (n->kind() == prim::CallFunction) {
+    auto func_node = n->inputs()[0]->node();
+    if (func_node->kind() == prim::Constant) {
+      is_quantizable |= std::find(call_funcs.begin(), call_funcs.end(), func_node->s(attr::name)) != call_funcs.end();
+    }
+  }
+  return is_quantizable;
+}
+
+bool valueNeedsToBeObserved(Value* v) {
+  if (!v->type()->isSubtypeOf(TensorType::get())) {
+    return false;
+  }
+  // Check whether producer is quantizable
+  if (nodeQuantizable(v->node())) {
+    return true;
+  }
+  // Check whether user is quantizable
+  for (const auto& use: v->uses()) {
+    if (nodeQuantizable(use.user)) {
+      return true;
+    }
+  }
+  return false;
 }
 
 Node* traverseToQuantNode(Node* dq) {
@@ -209,14 +241,17 @@ void InsertObserversImpl(
   // point is the beginning of graph node. This also safe guards against
   // observing a potentially mutated value due to some in-place operation
   Value* self = graph->inputs()[0];
+  std::unordered_set<Value*> values_observed;
   for (size_t idx = 1; idx < method.num_inputs(); ++idx) {
     auto& v = graph->inputs()[idx];
-    if (v->type()->isSubtypeOf(TensorType::get()) &&
-        values_to_skip.count(v) == 0) {
+    if (valueNeedsToBeObserved(v) &&
+        values_to_skip.count(v) == 0 &&
+        values_observed.count(v) == 0) {
       if (module_qconfig_map.count(module.module_object()) == 0) {
         // the module is added by us, it's an observer module
         continue;
       }
+      values_observed.emplace(v);
       auto qconfig = module_qconfig_map.at(module.module_object());
       if (qconfig) {
         auto observer_node =
@@ -235,15 +270,18 @@ void InsertObserversImpl(
     for (Node* n : b->nodes()) {
       // Skip nodes that we don't need to observe, e.g. 'prim::Constant' or
       // observer nodes
-      if (!outputsNeedToBeObserved(n) || observer_for_input.count(n) != 0) {
+      if (observer_for_input.count(n) != 0) {
         continue;
       }
 
       // Record all outputs in the values_to_observe - we'll later add observers
       // for all values from it.
       for (Value* v : n->outputs()) {
-        if (values_to_skip.count(v) == 0) {
+        if (valueNeedsToBeObserved(v) &&
+            values_to_skip.count(v) == 0 &&
+            values_observed.count(v) == 0) {
           values_to_observe.push_back(v);
+          values_observed.emplace(v);
         }
         if (v->node()->kind() == prim::CallMethod) {
           // If we find a call to a method of a child module,
@@ -277,9 +315,6 @@ void InsertObserversImpl(
 
   // Actually add observer nodes.
   for (Value* v : values_to_observe) {
-    if (!v->type()->isSubtypeOf(TensorType::get())) {
-      continue;
-    }
     // Skip inserting observer for bias
     if (v->node()->kind() == prim::GetAttr &&
         v->node()->s(c10::attr::name) == "bias") {
