Minor fix and turn off fold_convbn (#27403)

jerryzh168 · facebook-github-bot · commit 5ac3df7712a7 · 2019-11-01T12:15:38.000-07:00
Summary: Pull Request resolved: #27403 In fold_convbn pass, we need to recompute the parameter(weight, bias) for conv, update the attribute of conv and update the access of bias in conv because if the original conv have no bias, the `self.bias` access will be inline and replaced by Constant node `None = prim::Constant()`, we need to update this to use `GetAttr[name="bias"]` to make this work. But there is also some work going on the handle constants, so we'll fix this pass after that is done. Test Plan: . Imported from OSS Differential Revision: D18182918 fbshipit-source-id: bba510bc41ab58e0eb76f7b77335b6e3ffe2862d
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -1301,6 +1301,11 @@ def test_quant_fusion(self):
             FileCheck().run(input_str, graph)
 
     @_tmp_donotuse_dont_inline_everything
+    @unittest.skip("Temporarily turn off fold_convbn tests until \
+    constants are handled properly, this test should not be passing \
+    because bias is not handled properly, the reason is passes is because the \
+    parameters of bn are initialized to default values and the recomputed bias \
+    for conv is zero, which is equivalent to no bias")
     def test_foldbn_trivial(self):
         # Test trivial case
         class TestModule(torch.nn.Module):
@@ -1338,6 +1343,8 @@ def forward(self, x):
         self.assertAlmostEqual(eager(x), scripted(x), delta=1e-5)
 
     @_tmp_donotuse_dont_inline_everything
+    @unittest.skip("Temporarily turn off fold_convbn tests until \
+    constants are handled properly")
     def test_foldbn_trivial_nobias(self):
         # Test trivial case
         class TestModule(torch.nn.Module):
@@ -1375,6 +1382,8 @@ def forward(self, x):
         self.assertAlmostEqual(eager(x), scripted(x), delta=1e-5)
 
     @_tmp_donotuse_dont_inline_everything
+    @unittest.skip("Temporarily turn off fold_convbn tests until \
+    constants are handled properly")
     def test_foldbn_in_submodule(self):
         # Test that we find Conv-BN patterns in submodules
         class SubModule(torch.nn.Module):
diff --git a/torch/csrc/jit/passes/quantization.cpp b/torch/csrc/jit/passes/quantization.cpp
@@ -883,8 +883,12 @@ graph(%self, %x):
       GRAPH_UPDATE("Deleting ", *matched_bn);
 
       auto new_w_b = computeUpdatedConvWeightAndBias(params);
-      params.conv_w.set_data(std::get<0>(new_w_b));
-      params.conv_b.set_data(std::get<1>(new_w_b));
+      conv_submodule.set_parameter("weight", std::get<0>(new_w_b));
+      if (conv_submodule.find_parameter("bias")) {
+        conv_submodule.set_parameter("bias", std::get<1>(new_w_b));
+      } else {
+        conv_submodule.register_parameter("bias", std::get<1>(new_w_b), false);
+      }
     }
 
     // Perform planned rewritings
diff --git a/torch/quantization/_quantize_script.py b/torch/quantization/_quantize_script.py
@@ -132,7 +132,11 @@ def quantize_script(model, qconfig_dict, run_fn, run_args, inplace=False):
     if not inplace:
         model = model.copy()
     scripted_qconfig_dict = {k: script_qconfig(v) for k, v in qconfig_dict.items()}
-    torch._C._jit_pass_fold_convbn(model._c)
+    # We are not going to run fold_convbn pass right now
+    # since it is not able to work correctly, we will
+    # revisit after constants is properly handled in
+    # JIT
+    # torch._C._jit_pass_fold_convbn(model._c)
     prepare_script(model, scripted_qconfig_dict, True)
     run_fn(model._c._get_method('forward'), *run_args)
     # When we mutating graph we didn't create a new ClassType
