hook torch.jit.trace (#95)

Author: chunyuan-w

intel_pytorch_extension_py/ops/__init__.py

@@ -5,4 +5,4 @@
 from .reshape import *
 from .mlp import * 
 from .linear_fuse_relu import *
-from .jit_script import *
+from .jit import *

…l_pytorch_extension_py/ops/jit_script.py intel_pytorch_extension_py/ops/jit.pyintel_pytorch_extension_py/ops/jit_script.py renamed to intel_pytorch_extension_py/ops/jit.py intel_pytorch_extension_py/ops/jit_script.py renamed to intel_pytorch_extension_py/ops/jit.py

@@ -7,6 +7,7 @@
 torch._C._jit_set_profiling_executor(False)
 
 orig_script = torch.jit.script
+orig_trace = torch.jit.trace
 
 def script_(obj, optimize=None, _frames_up=0, _rcb=None):
     torch.jit.script = orig_script
@@ -20,8 +21,21 @@ def script_(obj, optimize=None, _frames_up=0, _rcb=None):
         ipex.enable_auto_mix_precision(None)
         jit_m = wrap_cpp_module(torch._C._jit_pass_fold_convbn(jit_m._c))
         ipex.enable_auto_mix_precision(orig_mixed_type)
+    return jit_m
+
+def trace_(func, example_inputs, *args, **kwargs):
+    # Disable mix precision. torch.jit.trace will check the traced output 
+    # against what is expected. Since mix precision will lead to 
+    # loss of accuracy, this will raise warning during torch.jit.trace
+    orig_mixed_type = ipex.get_auto_mix_precision()
+    ipex.enable_auto_mix_precision(None)
+    jit_m = orig_trace(func, example_inputs, *args, **kwargs)
 
+    if core.get_jit_opt():        
+        jit_m = wrap_cpp_module(torch._C._jit_pass_fold_convbn(jit_m._c))
+    ipex.enable_auto_mix_precision(orig_mixed_type)
     return jit_m
 
 
 torch.jit.script = script_
+torch.jit.trace = trace_

tests/cpu/test_jit.py

@@ -152,7 +152,7 @@ def forward(self, x):
 
 class Tester(TestCase):
 
-    def _test_output(self, model, x, kind=None):
+    def _test_output(self, model, x, kind_in_graph=None, kind_not_in_graph=None):
         modelName = model.__class__.__name__
         core.disable_jit_opt()
         core.disable_mix_bf16_fp32()
@@ -164,180 +164,207 @@ def _test_output(self, model, x, kind=None):
 
         script_model = torch.jit.script(model)
         script_model.eval()
+
+        trace_model = torch.jit.trace(model, x)
+        trace_model.eval()
         with torch.no_grad():
             sresult = script_model(x)
+            tresult = trace_model(x)
 
         self.assertEqual(result, sresult)
+        self.assertEqual(result, tresult)
 
         core.enable_jit_opt()
-        fused_model = torch.jit.script(model)
+        script_fused_model = torch.jit.script(model)
+        trace_fused_model = torch.jit.trace(model, x)
         with torch.no_grad():
             # conv relu fusion, conv sum fusion or conv sum relu fusion
-            graph =  fused_model.graph_for(x)
-            # print(graph)
-            fresult = fused_model(x)
+            script_graph =  script_fused_model.graph_for(x)
+            fused_sresult = script_fused_model(x)
 
-        # print(result)
-        # print(sresult)
-        # print(fresult)
+            trace_graph = trace_fused_model.graph_for(x)
+            fused_tresult = trace_fused_model(x)
 
-        self.assertEqual(result, fresult)
+        self.assertEqual(result, fused_sresult)
+        self.assertEqual(result, fused_tresult)
 
         # check if the fused node exists in the graph
-        if kind is not None:
-            self.assertTrue(any(n.kind() == kind for n in graph.nodes()))
+        if kind_in_graph is not None:
+            self.assertTrue(any(n.kind() == kind_in_graph for n in script_graph.nodes()))
+            self.assertTrue(any(n.kind() == kind_in_graph for n in trace_graph.nodes()))
+        
+        # check if certain node does not exist in the graph
+        if kind_not_in_graph is not None:
+            self.assertTrue(all(n.kind() != kind_not_in_graph for n in script_graph.nodes()))
+            self.assertTrue(all(n.kind() != kind_not_in_graph for n in trace_graph.nodes()))
+
 
-    def _test_output_bf16(self, model, x, kind=None, prec=None):
+    def _test_output_bf16(self, model, x, kind_in_graph=None, kind_not_in_graph=None, prec=None):
         modelName = model.__class__.__name__
 
         core.enable_auto_dnnl()
         core.enable_jit_opt()
-        core.disable_mix_bf16_fp32()
-
+        core.enable_mix_bf16_fp32()
+        
         model = model.to(ipex.DEVICE).eval()
         x = x.to(ipex.DEVICE)
         x2 = x.clone()
+        x3 = x.clone()
+        
+        script_fused_model = torch.jit.script(copy.deepcopy(model))
+        trace_fused_model = torch.jit.trace(copy.deepcopy(model), x3)
 
-        fused_model = torch.jit.script(copy.deepcopy(model))
-
-        # bn folding, removing it after solve some issue, using mix_preci? to check
-        core.disable_auto_dnnl()
-        fused_model = wrap_cpp_module(torch._C._jit_pass_fold_convbn(fused_model._c))
-        core.enable_auto_dnnl()
-
-        core.enable_mix_bf16_fp32()
 
         with torch.no_grad():
             # bf16, native path
             result = model(x)
-            # bf16, jit path
-            graph =  fused_model.graph_for(x2)
-            # print(graph)
-            fresult = fused_model(x2)
-
-        #print(result)
-        #print(fresult)
-
-        self.assertEqual(fresult, result, prec=prec)
+            # bf16, jit script path
+            script_graph =  script_fused_model.graph_for(x2)
+            fused_sresult = script_fused_model(x2)
+            # bf 16, jit trace path
+            trace_graph = trace_fused_model.graph_for(x3)
+            fused_tresult = trace_fused_model(x3)
+        
+        # disable mix_bf16_fp32 when the calculation is done 
+        # to avoid affecting other scripts
+        core.disable_mix_bf16_fp32()
+        
+        self.assertEqual(fused_sresult, result, prec=prec)
+        self.assertEqual(fused_tresult, result, prec=prec)
 
         # check if the fused node exists in the graph
-        if kind is not None:
-            self.assertTrue(any(n.kind() == kind for n in graph.nodes()))
-
+        if kind_in_graph is not None:
+            self.assertTrue(any(n.kind() == kind_in_graph for n in script_graph.nodes()))
+            self.assertTrue(any(n.kind() == kind_in_graph for n in trace_graph.nodes()))
+        
+        # check if certain node does not exist in the graph
+        if kind_not_in_graph is not None:
+            self.assertTrue(all(n.kind() != kind_not_in_graph for n in script_graph.nodes()))
+            self.assertTrue(all(n.kind() != kind_not_in_graph for n in trace_graph.nodes()))
+    
 
     def test_output_conv_bn_2d(self):
         self._test_output(
             ConvBatchNorm_Fixed(2, 3, 32, kernel_size=3, stride=1),
-            torch.randn(32, 3, 224, 224),
-            kind="aten::conv2d")
+            torch.randn(32, 3, 64, 64),
+            kind_in_graph="aten::conv2d",
+            kind_not_in_graph="aten::batch_norm",)
         self._test_output_bf16(
             ConvBatchNorm_Fixed(2, 3, 32, kernel_size=3, stride=1),
-            torch.randn(32, 3, 224, 224),
-            kind="aten::conv2d",
+            torch.randn(32, 3, 64, 64),
+            kind_in_graph="aten::conv2d",
+            kind_not_in_graph="aten::batch_norm",
             prec=0.02)
 
 
     def test_output_conv_bn_3d(self):
         self._test_output(
             ConvBatchNorm_Fixed(3, 3, 32, kernel_size=3, stride=1),
-            torch.randn(32, 3, 112, 112, 112),
-            kind="aten::conv3d")
+            torch.randn(32, 3, 32, 32, 32),
+            kind_in_graph="aten::conv3d",
+            kind_not_in_graph="aten::batch_norm",)
         self._test_output_bf16(
             ConvBatchNorm_Fixed(3, 3, 32, kernel_size=3, stride=1),
-            torch.randn(32, 3, 112, 112, 112),
-            kind="aten::conv3d",
+            torch.randn(32, 3, 32, 32, 32),
+            kind_in_graph="aten::conv3d",
+            kind_not_in_graph="aten::batch_norm",
             prec=0.02)
 
 
     def test_output_conv_relu_2d(self):
         self._test_output(
             ConvRelu_Fixed(2, 3, 32, kernel_size=3, stride=1),
-            torch.randn(32, 3, 224, 224),
-            kind="ipex::conv2d_relu")
+            torch.randn(32, 3, 64, 64),
+            kind_in_graph="ipex::conv2d_relu")
         self._test_output_bf16(
             ConvRelu_Fixed(2, 3, 32, kernel_size=3, stride=1),
-            torch.randn(32, 3, 224, 224),
-            kind="ipex::conv2d_relu")
+            torch.randn(32, 3, 64, 64),
+            kind_in_graph="ipex::conv2d_relu")
 
 
     def test_output_conv_relu_3d(self):
         self._test_output(
             ConvRelu_Fixed(3, 3, 32, kernel_size=3, stride=1),
-            torch.randn(32, 3, 112, 112, 112),
-            kind="ipex::conv3d_relu")
+            torch.randn(32, 3, 32, 32, 32),
+            kind_in_graph="ipex::conv3d_relu")
         self._test_output_bf16(
             ConvRelu_Fixed(3, 3, 32, kernel_size=3, stride=1),
-            torch.randn(32, 3, 112, 112, 112),
-            kind="ipex::conv3d_relu")
+            torch.randn(32, 3, 32, 32, 32),
+            kind_in_graph="ipex::conv3d_relu")
 
 
     def test_output_conv_sum_2d(self):
         self._test_output(
             ConvSum(2, 3, 32, kernel_size=3, stride=1),
-            torch.randn(32, 3, 224, 224),
-            kind="ipex::conv2d_sum")
+            torch.randn(32, 3, 64, 64),
+            kind_in_graph="ipex::conv2d_sum")
         self._test_output_bf16(
             ConvSum(2, 3, 32, kernel_size=3, stride=1),
-            torch.randn(32, 3, 224, 224),
-            kind="ipex::conv2d_sum",
+            torch.randn(32, 3, 64, 64),
+            kind_in_graph="ipex::conv2d_sum",
             prec=0.04)
 
 
     def test_output_conv_sum_3d(self):
         self._test_output(
             ConvSum(3, 3, 32, kernel_size=3, stride=1),
-            torch.randn(32, 3, 112, 112, 112),
-            kind="ipex::conv3d_sum")
+            torch.randn(32, 3, 32, 32, 32),
+            kind_in_graph="ipex::conv3d_sum")
         self._test_output_bf16(
             ConvSum(3, 3, 32, kernel_size=3, stride=1),
-            torch.randn(32, 3, 112, 112, 112),
-            kind="ipex::conv3d_sum",
+            torch.randn(32, 3, 32, 32, 32),
+            kind_in_graph="ipex::conv3d_sum",
             prec=0.04)
 
 
     def test_output_cascaded_conv_bn_sum_relu_2d(self):
         self._test_output(
             CascadedConvBnSumRelu(2, 3, 64, 32, kernel_size=3, stride=1),
-            torch.rand(32, 3, 224, 224),
-            kind="ipex::conv2d_sum_relu")
+            torch.rand(32, 3, 64, 64),
+            kind_in_graph="ipex::conv2d_sum_relu",
+            kind_not_in_graph="aten::batch_norm")
         self._test_output_bf16(
             CascadedConvBnSumRelu(2, 3, 64, 32, kernel_size=3, stride=1),
-            torch.rand(32, 3, 224, 224),
-            kind="ipex::conv2d_sum_relu",
+            torch.rand(32, 3, 64, 64),
+            kind_in_graph="ipex::conv2d_sum_relu",
+            kind_not_in_graph="aten::batch_norm",
             prec=0.02)
 
 
     def test_output_cascaded_conv_bn_sum_relu_3d(self):
         self._test_output(
             CascadedConvBnSumRelu(3, 3, 64, 32, kernel_size=3, stride=1),
-            torch.rand(32, 3, 112, 112, 112),
-            kind="ipex::conv3d_sum_relu")
+            torch.rand(32, 3, 32, 32, 32),
+            kind_in_graph="ipex::conv3d_sum_relu",
+            kind_not_in_graph="aten::batch_norm",)
         self._test_output_bf16(
             CascadedConvBnSumRelu(3, 3, 64, 32, kernel_size=3, stride=1),
-            torch.rand(32, 3, 112, 112, 112),
-            kind="ipex::conv3d_sum_relu",
+            torch.rand(32, 3, 32, 32, 32),
+            kind_in_graph="ipex::conv3d_sum_relu",
+            kind_not_in_graph="aten::batch_norm",
             prec=0.02)
 
 
     def test_output_linear_relu(self):
         self._test_output(
             LinearRelu(3, 32, bias=True),
             torch.rand(32, 3),
-            kind="ipex::linear_relu")
+            kind_in_graph="ipex::linear_relu")
         self._test_output_bf16(
             LinearRelu(3, 32, bias=True),
             torch.rand(32, 3),
-            kind="ipex::linear_relu")
+            kind_in_graph="ipex::linear_relu")
         self._test_output(
             LinearRelu(3, 32, bias=False),
             torch.rand(32, 3),
-            kind="ipex::linear_relu")
+            kind_in_graph="ipex::linear_relu")
         self._test_output_bf16(
             LinearRelu(3, 32, bias=False),
             torch.rand(32, 3),
-            kind="ipex::linear_relu")
+            kind_in_graph="ipex::linear_relu")
 
 
 if __name__ == '__main__':
+    torch.manual_seed(2020)
     core.enable_auto_dnnl()
     test = unittest.main()