_jit_pass_fold_convbn wrapped with fuse_conv_bn_script (pytorch#40224)

Summary: Pull Request resolved: pytorch#40224

Test Plan: Imported from OSS

Differential Revision: D22117111

Pulled By: edmundw314

fbshipit-source-id: 9252674bd770ba6669d50090849d9f9bc13edaa3

Author: edmundw314

test/quantization/test_quantize_script.py

@@ -23,6 +23,7 @@
 from torch.quantization.quantize_script import prepare_dynamic_script
 from torch.quantization.quantize_script import convert_dynamic_script
 from torch.quantization.quantize_script import quantize_dynamic_script
+from torch.quantization.quantize_script import fuse_conv_bn_script
 
 # Testing utils
 from torch.testing._internal.common_quantization import test_only_eval_fn as _test_only_eval_fn
@@ -86,7 +87,7 @@ def forward(self, x):
                 .run(str(get_forward(scripted_or_traced._c).graph))
 
             # Run FoldConvBatchnorm2d pass.
-            scripted_or_traced = wrap_cpp_module(torch._C._jit_pass_fold_convbn(scripted_or_traced._c))
+            scripted_or_traced = fuse_conv_bn_script(scripted_or_traced)
 
             # Check that after the pass one of the CallMethods is gone (supposedly,
             # the bn.forward).
@@ -130,7 +131,7 @@ def forward(self, x):
                 .run(str(get_forward_graph(scripted_or_traced._c)))
 
             # Run FoldConvBatchnorm2d pass.
-            scripted_or_traced = wrap_cpp_module(torch._C._jit_pass_fold_convbn(scripted_or_traced._c))
+            scripted_or_traced = fuse_conv_bn_script(scripted_or_traced)
 
             # Check that after the pass one of the CallMethods is gone (supposedly,
             # the bn.forward).
@@ -176,7 +177,7 @@ def forward(self, x):
             FileCheck().check_count("prim::CallMethod[name=\"forward\"]", 2, exactly=True) \
                 .run(str(get_forward_graph(scripted_or_traced.sub._c)))
 
-            scripted_or_traced = wrap_cpp_module(torch._C._jit_pass_fold_convbn(scripted_or_traced._c))
+            scripted_or_traced = fuse_conv_bn_script(scripted_or_traced)
 
             FileCheck().check_count("prim::CallMethod[name=\"forward\"]", 1, exactly=True) \
                 .run(str(get_forward_graph(scripted_or_traced.sub._c)))
@@ -227,7 +228,7 @@ def forward(self, x):
             FileCheck().check_count("prim::CallMethod[name=\"forward\"]", 2, exactly=True) \
                 .run(str(get_forward_graph(scripted_or_traced.sub._c)))
 
-            scripted_or_traced = wrap_cpp_module(torch._C._jit_pass_fold_convbn(scripted_or_traced._c))
+            scripted_or_traced = fuse_conv_bn_script(scripted_or_traced)
 
             FileCheck().check_count("prim::CallMethod[name=\"forward\"]", 2, exactly=True) \
                 .run(str(get_forward_graph(scripted_or_traced.sub._c)))
@@ -268,7 +269,7 @@ def forward(self, x):
                 else:
                     scripted_or_traced = torch.jit.script(eager).copy()
                 torch._C._jit_pass_dedup_module_uses(scripted_or_traced ._c)
-                folded = wrap_cpp_module(torch._C._jit_pass_fold_convbn(scripted_or_traced ._c))
+                folded = fuse_conv_bn_script(scripted_or_traced)
                 x = torch.rand(1, 5, 6, 6)
                 self.assertEqual(eager(x), scripted_or_traced(x))
 
@@ -321,7 +322,7 @@ def forward(self, x):
             FileCheck().check_count("prim::CallMethod[name=\"forward\"]", num_layers * 2, exactly=True) \
                 .run(str(get_forward_graph(scripted_or_traced.sub.layers._c)))
 
-            scripted_or_traced = wrap_cpp_module(torch._C._jit_pass_fold_convbn(scripted_or_traced._c))
+            scripted_or_traced = fuse_conv_bn_script(scripted_or_traced)
 
             FileCheck().check_count("prim::CallMethod[name=\"forward\"]", num_layers, exactly=True) \
                 .run(str(get_forward_graph(scripted_or_traced.sub.layers._c)))

torch/quantization/quantize_script.py

@@ -27,14 +27,17 @@ def script_qconfig(qconfig):
 def script_qconfig_dict(qconfig_dict):
     return {k: script_qconfig(v) if v else None for k, v in qconfig_dict.items()}
 
+def fuse_conv_bn_script(model):
+    return torch.jit._recursive.wrap_cpp_module(torch._C._jit_pass_fold_convbn(model._c))
+
 def _prepare_script(model, qconfig_dict, inplace=False, quant_type=QuantType.STATIC):
     assert not inplace, "The inplace support is still in development"
     _check_is_script_module(model)
     _check_forward_method(model)
     if not all(isinstance(x, str) for x in qconfig_dict.keys()):
         raise ValueError('qconfig_dict should only contain names(str) as keys.')
     scripted_qconfig_dict = script_qconfig_dict(qconfig_dict)
-    model = wrap_cpp_module(torch._C._jit_pass_fold_convbn(model._c))
+    model = fuse_conv_bn_script(model)
     return wrap_cpp_module(torch._C._jit_pass_insert_observers(model._c,
                                                                'forward',
                                                                scripted_qconfig_dict,