Torchscript backend (#17132)

* Added torch script backend

* Added ts_backend to pytorch layer tests

* Added use_ts_backend fixture to the test suite to activate the
torchscript backend

* Fixed failing test_dict layer test

* Added USE_TS_BACKEND as an env variable

* Removed use_ts_backend fixture

* Added more tests for ts backend

* Added more information in the comments about usage

* Removed convolution3d test from precommit_ts_backend

* Added some torchscript backend tests to ci

* Removed tests from CI as torch.compile doesn't support 3.11 currently

* Fixed linter issues

* Addressed PR comments and linter issues

.ci/azure/linux.yml

@@ -105,7 +105,7 @@ jobs:
   steps:
   - task: UsePythonVersion@0
     inputs:
-      versionSpec: '$(OV_PYTHON_VERSION)' # Setting only major & minor version will download latest release from GH repo example 3.10 will be 3.10.10. 
+      versionSpec: '$(OV_PYTHON_VERSION)' # Setting only major & minor version will download latest release from GH repo example 3.10 will be 3.10.10.
       addToPath: true
       disableDownloadFromRegistry: false
       architecture: 'x64'
@@ -367,7 +367,7 @@ jobs:
     displayName: 'Build cpp samples - gcc'
 
   - script: $(SAMPLES_INSTALL_DIR)/cpp/build_samples.sh -b $(BUILD_DIR)/cpp_samples_clang
-    env: 
+    env:
       CC: clang
       CXX: clang++
     displayName: 'Build cpp samples - clang'

src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/backend.py

@@ -0,0 +1,98 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+import logging
+import os
+import torch
+from torch._dynamo.backends.common import fake_tensor_unsupported
+from torch._dynamo.backends.registry import register_backend
+from torch._inductor.compile_fx import compile_fx
+
+from openvino.frontend import FrontEndManager
+from openvino.runtime import Core, Type, PartialShape
+from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+
+log = logging.getLogger(__name__)
+
+"""
+    This is a preview feature in OpenVINO. Torchscript backend
+    enables users to compile PyTorch models using torch.compile
+    with OpenVINO as a target backend in PyTorch applications
+
+    Sample usage:
+    This sample code loads resnet50 torchvision model and compiles it using torch dynamo.
+    We can then use this model for inference. We only need to add two lines of code to
+    the Pytorch applications which are marked in the code below
+
+    1) import openvino.frontend.pytorch.torchdynamo.backend
+    model = torchvision.models.resnet50()
+    2) model = torch.compile(model, backend="openvino")
+"""
+
+
+@register_backend
+@fake_tensor_unsupported
+def openvino(subgraph, example_inputs):
+    return ts_openvino(subgraph, example_inputs)
+
+
+def ts_openvino(subgraph, example_inputs):
+    try:
+        model = torch.jit.script(subgraph)
+        model.eval()
+        fr_model = torch.jit.freeze(model)
+
+        core = Core()
+        fe_manager = FrontEndManager()
+        fe = fe_manager.load_by_framework('pytorch')
+        dtype_mapping = {
+            torch.float64: Type.f64,
+            torch.float32: Type.f32,
+            torch.float16: Type.f16,
+            torch.int64: Type.i64,
+            torch.int32: Type.i32,
+            torch.uint8: Type.u8,
+            torch.int8: Type.i8,
+            torch.bool: Type.boolean,
+        }
+        decoder = TorchScriptPythonDecoder(fr_model)
+
+        # TODO: Use convert_model instead when mo --convert_model api becomes a part of OV runtime
+        im = fe.load(decoder)
+        om = fe.convert(im)
+
+        for idx, input_data in enumerate(example_inputs):
+            om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype])
+            om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape)))
+        om.validate_nodes_and_infer_types()
+
+        device = "CPU"
+        if (os.getenv("OPENVINO_TS_BACKEND_DEVICE") is not None):
+            device = os.getenv("OPENVINO_TS_BACKEND_DEVICE")
+            assert device in core.available_devices, "Specified device " + device + " is not in the list of OpenVINO Available Devices"
+
+        compiled_model = core.compile_model(om, device)
+
+        def _call(*args):
+            if not hasattr(_call, "execute_on_ov"):
+                _call.execute_on_ov = True
+            execute_on_ov = getattr(_call, "execute_on_ov")
+            if execute_on_ov:
+                ov_inputs = [a.detach().cpu().numpy() for a in args]
+                try:
+                    res = compiled_model(ov_inputs)
+                except Exception as e:
+                    log.debug(f"Failed in OpenVINO execution: {e}")
+                    _call.execute_on_ov = False
+                    return subgraph.forward(*args)
+                result = [torch.from_numpy(res[out]) for out in compiled_model.outputs]
+                return result
+            else:
+                return subgraph.forward(*args)
+        return _call
+    except Exception as e:
+        log.debug(f"Failed in compilation: {e}")
+        return compile_fx(subgraph, example_inputs)

tests/layer_tests/pytest.ini

@@ -2,4 +2,5 @@
 markers =
     nightly
     precommit
+    precommit_ts_backend
     timeout

tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py

@@ -4,6 +4,7 @@
 import itertools
 import warnings
 from copy import deepcopy
+import os
 
 import numpy as np
 from common.constants import test_device, test_precision
@@ -11,6 +12,8 @@ from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
 
 from openvino.frontend import FrontEndManager
 from openvino.runtime import Core, Type, PartialShape
+import torch
+import openvino.frontend.pytorch.torchdynamo.backend
 
 
 class PytorchLayerTest:
@@ -46,104 +49,89 @@ class PytorchLayerTest:
             inputs = self._prepare_input(**kwargs['kwargs_to_prepare_input'])
         else:
             inputs = self._prepare_input()
-        with torch.no_grad():
-            model.eval()
-            torch_inputs = [torch.from_numpy(inp) if isinstance(
-                inp, np.ndarray) else inp for inp in inputs]
-            trace_model = kwargs.get('trace_model', False)
-            freeze_model = kwargs.get('freeze_model', True)
-            model, converted_model = self.convert_directly_via_frontend(model, torch_inputs, trace_model, dynamic_shapes, inputs, freeze_model)
-            graph = model.inlined_graph
 
-            if kind is not None and not isinstance(kind, (tuple, list)):
-                kind = [kind]
-            if kind is not None:
-                for op in kind:
-                    assert self._check_kind_exist(
-                        graph, op), f"Operation {op} type doesn't exist in provided graph"
-
-        # OV infer:
-        core = Core()
-        compiled = core.compile_model(converted_model, ie_device)
-        infer_res = compiled(deepcopy(inputs))
-
-        if hasattr(self, 'skip_framework') and self.skip_framework:
-            warnings.warn('Framework is skipped')
-            return
-
-        # Framework infer:
-        fw_res = model(*deepcopy(torch_inputs))
-
-        if not isinstance(fw_res, (tuple)):
-            fw_res = (fw_res,)
-
-        output_list = list(infer_res.values())
-
-        flatten_fw_res = []
-
-        def flattenize_dict_outputs(res):
-            if isinstance(res, dict):
-                return flattenize_outputs(res.values())
-            
-        def flattenize_outputs(res):
-            results = []
-            for res_item in res:
-                # if None is at output we skip it
-                if res_item is None:
-                    continue
-                # If input is list or tuple flattenize it
-                if isinstance(res_item, (list, tuple)):
-                    decomposed_res = flattenize_outputs(res_item)
-                    results.extend(decomposed_res)
-                    continue
-                if isinstance(res_item, dict):
-                    decomposed_res = flattenize_dict_outputs(res_item)
-                    results.extend(decomposed_res)
-                    continue
-                results.append(res_item)
-            return results
-
-        flatten_fw_res = flattenize_outputs(fw_res)
-
-        assert len(flatten_fw_res) == len(
-            output_list), f'number of outputs not equal, {len(flatten_fw_res)} != {len(output_list)}'
-        # check if results dtypes match
-        for fw_tensor, ov_tensor in zip(flatten_fw_res, output_list):
-            if not isinstance(fw_tensor, torch.Tensor):
-                fw_type = torch.tensor(fw_tensor).numpy().dtype
-                ov_type = ov_tensor.dtype
-                if fw_type in [np.int32, np.int64] and ov_type in [np.int32, np.int64]:
-                    # do not differentiate between int32 and int64
-                    continue
-                assert ov_type == fw_type, f"dtype validation failed: {ov_type} != {fw_type}"
-                continue
-            assert torch.tensor(np.array(
-                ov_tensor)).dtype == fw_tensor.dtype, f"dtype validation failed: {torch.tensor(np.array(ov_tensor)).dtype} != {fw_tensor.dtype}"
+        torch_inputs = [torch.from_numpy(inp) if isinstance(
+            inp, np.ndarray) else inp for inp in inputs]
 
         if 'custom_eps' in kwargs and kwargs['custom_eps'] is not None:
             custom_eps = kwargs['custom_eps']
         else:
             custom_eps = 1e-4
 
-        # Compare Ie results with Framework results
-        fw_eps = custom_eps if precision == 'FP32' else 5e-2
-        is_ok = True
-        for i in range(len(infer_res)):
-            cur_fw_res = flatten_fw_res[i].to(memory_format=torch.contiguous_format).numpy(
-            ) if isinstance(flatten_fw_res[i], torch.Tensor) else flatten_fw_res[i]
-            cur_ov_res = infer_res[compiled.output(i)]
-            print(f"fw_re: {cur_fw_res};\n ov_res: {cur_ov_res}")
-            if not np.allclose(cur_ov_res, cur_fw_res,
-                               atol=fw_eps,
-                               rtol=fw_eps, equal_nan=True):
-                is_ok = False
-                print("Max diff is {}".format(
-                    np.array(
-                        abs(cur_ov_res - cur_fw_res)).max()))
-            else:
-                print("Accuracy validation successful!\n")
-                print("absolute eps: {}, relative eps: {}".format(fw_eps, fw_eps))
-        assert is_ok, "Accuracy validation failed"
+        def use_ts_backend():
+            return(os.environ.get('USE_TS_BACKEND', False))
+
+        if use_ts_backend():
+            self.ts_backend_test(model, torch_inputs, custom_eps)
+        else:
+            with torch.no_grad():
+                model.eval()
+                trace_model = kwargs.get('trace_model', False)
+                freeze_model = kwargs.get('freeze_model', True)
+                model, converted_model = self.convert_directly_via_frontend(model, torch_inputs, trace_model, dynamic_shapes, inputs, freeze_model)
+                graph = model.inlined_graph
+
+                if kind is not None and not isinstance(kind, (tuple, list)):
+                    kind = [kind]
+                if kind is not None:
+                    for op in kind:
+                        assert self._check_kind_exist(
+                            graph, op), f"Operation {op} type doesn't exist in provided graph"
+            # OV infer:
+            core = Core()
+            compiled = core.compile_model(converted_model, ie_device)
+            infer_res = compiled(deepcopy(inputs))
+
+            if hasattr(self, 'skip_framework') and self.skip_framework:
+                warnings.warn('Framework is skipped')
+                return
+
+            # Framework infer:
+            fw_res = model(*deepcopy(torch_inputs))
+
+            if not isinstance(fw_res, (tuple)):
+                fw_res = (fw_res,)
+
+            output_list = list(infer_res.values())
+
+            flatten_fw_res = []
+
+            flatten_fw_res = flattenize_outputs(fw_res)
+
+            assert len(flatten_fw_res) == len(
+                output_list), f'number of outputs not equal, {len(flatten_fw_res)} != {len(output_list)}'
+            # check if results dtypes match
+            for fw_tensor, ov_tensor in zip(flatten_fw_res, output_list):
+                if not isinstance(fw_tensor, torch.Tensor):
+                    fw_type = torch.tensor(fw_tensor).numpy().dtype
+                    ov_type = ov_tensor.dtype
+                    if fw_type in [np.int32, np.int64] and ov_type in [np.int32, np.int64]:
+                        # do not differentiate between int32 and int64
+                        continue
+                    assert ov_type == fw_type, f"dtype validation failed: {ov_type} != {fw_type}"
+                    continue
+                assert torch.tensor(np.array(
+                    ov_tensor)).dtype == fw_tensor.dtype, f"dtype validation failed: {torch.tensor(np.array(ov_tensor)).dtype} != {fw_tensor.dtype}"
+
+            # Compare Ie results with Framework results
+            fw_eps = custom_eps if precision == 'FP32' else 5e-2
+            is_ok = True
+            for i in range(len(infer_res)):
+                cur_fw_res = flatten_fw_res[i].to(memory_format=torch.contiguous_format).numpy(
+                ) if isinstance(flatten_fw_res[i], torch.Tensor) else flatten_fw_res[i]
+                cur_ov_res = infer_res[compiled.output(i)]
+                print(f"fw_re: {cur_fw_res};\n ov_res: {cur_ov_res}")
+                if not np.allclose(cur_ov_res, cur_fw_res,
+                                atol=fw_eps,
+                                rtol=fw_eps, equal_nan=True):
+                    is_ok = False
+                    print("Max diff is {}".format(
+                        np.array(
+                            abs(cur_ov_res - cur_fw_res)).max()))
+                else:
+                    print("Accuracy validation successful!\n")
+                    print("absolute eps: {}, relative eps: {}".format(fw_eps, fw_eps))
+            assert is_ok, "Accuracy validation failed"
 
     # Each model should specify inputs
     def _prepare_input(self):
@@ -204,6 +192,58 @@ class PytorchLayerTest:
         om.validate_nodes_and_infer_types()
         return om
 
+    def ts_backend_test(self, model, inputs, custom_eps):
+        torch._dynamo.reset()
+        with torch.no_grad():
+            model.eval()
+            fw_model = torch.compile(model)
+            ov_model = torch.compile(model, backend="openvino")
+        ov_res = ov_model(*inputs)
+        fw_res = fw_model(*inputs)
+
+        if not isinstance(fw_res, (tuple)):
+            fw_res = (fw_res,)
+
+        if not isinstance(ov_res, (tuple)):
+            ov_res = (ov_res,)
+
+        flatten_fw_res, flatten_ov_res = [], []
+        flatten_fw_res = flattenize_outputs(fw_res)
+        flatten_ov_res = flattenize_outputs(ov_res)
+
+        assert len(flatten_fw_res) == len(
+            flatten_ov_res
+        ), f'number of ouptuts are not equal, {len(flatten_fw_res)} != {len(flatten_ov_res)}'
+
+
+        # Check if output data types match
+        for fw_tensor, ov_tensor in zip(flatten_fw_res, flatten_ov_res):
+            if not isinstance(fw_tensor, torch.Tensor) and not isinstance(ov_tensor, torch.Tensor):
+                assert fw_tensor == ov_tensor
+                assert type(fw_tensor) == type(ov_tensor)
+                continue
+            assert fw_tensor.dtype == ov_tensor.dtype, f"dtype validation failed: {fw_tensor.dtype} != {ov_tensor.dtype}"
+
+        fw_eps = custom_eps
+        is_ok = True
+        for i in range(len(flatten_ov_res)):
+            cur_ov_res = flatten_ov_res[i]
+            cur_fw_res = flatten_fw_res[i]
+            if not torch.allclose(cur_fw_res, cur_ov_res,
+                                  atol=fw_eps, rtol=fw_eps,
+                                  equal_nan=True):
+                is_ok = False
+                print(
+                    "Max diff is {}".format(
+                        torch.max(torch.tensor(abs(cur_ov_res - cur_fw_res)))
+                    )
+                )
+            else:
+                print("Accuracy validation successful!\n")
+                print("absolute eps: {}, relative eps: {}".format(fw_eps, fw_eps))
+        assert is_ok, "Accuracy validation failed"
+
+
 
 def get_params(ie_device=None, precision=None):
     """
@@ -220,3 +260,27 @@ def get_params(ie_device=None, precision=None):
             continue
         test_args.append(element)
     return test_args
+
+
+def flattenize_dict_outputs(res):
+    if isinstance(res, dict):
+        return flattenize_outputs(res.values())
+
+
+def flattenize_outputs(res):
+    results = []
+    for res_item in res:
+        # if None is at output we skip it
+        if res_item is None:
+            continue
+        # If input is list or tuple flattenize it
+        if isinstance(res_item, (list, tuple)):
+            decomposed_res = flattenize_outputs(res_item)
+            results.extend(decomposed_res)
+            continue
+        if isinstance(res_item, dict):
+            decomposed_res = flattenize_dict_outputs(res_item)
+            results.extend(decomposed_res)
+            continue
+        results.append(res_item)
+    return results

tests/layer_tests/pytorch_tests/test_adaptive_avg_pool3d.py

@@ -32,6 +32,7 @@ class TestAdaptiveAvgPool3D(PytorchLayerTest):
 
     @pytest.mark.nightly
     @pytest.mark.precommit
+    @pytest.mark.precommit_ts_backend
     def test_adaptive_avg_pool3d(self, ie_device, precision, ir_version, input_tensor, output_size):
         self.input_tensor = input_tensor
         self._test(*self.create_model(output_size), ie_device, precision, ir_version)

tests/layer_tests/pytorch_tests/test_adaptive_max_pool_2d.py

@@ -46,6 +46,7 @@ class TestAdaptiveMaxPool2D(PytorchLayerTest):
     ]))
     @pytest.mark.nightly
     @pytest.mark.precommit
+    @pytest.mark.precommit_ts_backend
     def test_adaptive_max_pool2d(self, ie_device, precision, ir_version, input_tensor, output_size, return_indices):
         self.input_tensor = input_tensor
         self._test(*self.create_model(output_size, return_indices), ie_device, precision, ir_version)

tests/layer_tests/pytorch_tests/test_add.py

@@ -38,6 +38,7 @@ class TestAdd(PytorchLayerTest):
 
     @pytest.mark.nightly
     @pytest.mark.precommit
+    @pytest.mark.precommit_ts_backend
     @pytest.mark.parametrize("op_type", ["add", "add_"])
     def test_add(self, ie_device, precision, ir_version, alpha, input_rhs, op_type):
         self.input_rhs = input_rhs
@@ -98,6 +99,7 @@ class TestAddTypes(PytorchLayerTest):
                                                           ])
     @pytest.mark.nightly
     @pytest.mark.precommit
+    @pytest.mark.precommit_ts_backend
     def test_add_types(self, ie_device, precision, ir_version, lhs_type, lhs_shape, rhs_type, rhs_shape):
         self.lhs_type = lhs_type
         self.lhs_shape = lhs_shape
@@ -120,5 +122,6 @@ class TestAddLists(PytorchLayerTest):
 
     @pytest.mark.nightly
     @pytest.mark.precommit
+    @pytest.mark.precommit_ts_backend
     def test_add(self, ie_device, precision, ir_version):
         self._test(*self.create_model(), ie_device, precision, ir_version)

tests/layer_tests/pytorch_tests/test_batch_norm.py

@@ -34,7 +34,7 @@ class TestBatchNorm(PytorchLayerTest):
                 super(aten_batch_norm_train, self).__init__()
                 self.weight = torch.randn(6) if weights else None
                 self.bias = torch.randn(6) if bias else None
-                self.running_mean = torch.randn(6) if running_stats else None 
+                self.running_mean = torch.randn(6) if running_stats else None
                 self.running_var = torch.randn(6) if running_stats else None
                 self.eps = eps
 
@@ -57,6 +57,7 @@ class TestBatchNorm(PytorchLayerTest):
     ])
     @pytest.mark.nightly
     @pytest.mark.precommit
+    @pytest.mark.precommit_ts_backend
     def test_batch_norm(self, weights, bias, eps, train, running_stats, ie_device, precision, ir_version, kwargs_to_prepare_input):
         self._test(*self.create_model(weights, bias, eps, train, running_stats),
                    ie_device, precision, ir_version, kwargs_to_prepare_input=kwargs_to_prepare_input, dynamic_shapes=False, use_mo_convert=False)

tests/layer_tests/pytorch_tests/test_convolution.py

@@ -207,6 +207,7 @@ class TestConvolution(PytorchLayerTest):
     @pytest.mark.parametrize("underscore", [True, False])
     @pytest.mark.nightly
     @pytest.mark.precommit
+    @pytest.mark.precommit_ts_backend
     def test_convolution1d(self, params, bias, underscore, ie_device, precision, ir_version):
         self._test(*self.create_model(**params, bias=bias, underscore=underscore),
                    ie_device, precision, ir_version, dynamic_shapes=params['groups'] == 1,
@@ -217,6 +218,7 @@ class TestConvolution(PytorchLayerTest):
     @pytest.mark.parametrize("underscore", [True, False])
     @pytest.mark.nightly
     @pytest.mark.precommit
+    @pytest.mark.precommit_ts_backend
     def test_convolution2d(self, params, bias, underscore, ie_device, precision, ir_version):
         self._test(*self.create_model(**params, bias=bias, underscore=underscore),
                    ie_device, precision, ir_version, dynamic_shapes=params['groups'] == 1)

tests/layer_tests/pytorch_tests/test_relu.py

@@ -30,5 +30,6 @@ class TestRelu(PytorchLayerTest):
     @pytest.mark.parametrize("inplace", [False, True])
     @pytest.mark.nightly
     @pytest.mark.precommit
+    @pytest.mark.precommit_ts_backend
     def test_relu(self, inplace, ie_device, precision, ir_version):
         self._test(*self.create_model(inplace), ie_device, precision, ir_version)