From d018779adda8067952e020e9d946b11924d00fdb Mon Sep 17 00:00:00 2001 From: Maxim Vafin Date: Tue, 30 Jan 2024 12:50:18 +0100 Subject: [PATCH] [PT FE] torch.export support (#22397) * [PT FE] torch.export support * Apply code style * Fix build * Support torch <2.2 * Support more operations * Update tests/model_hub_tests/torch_tests/torch_utils.py * Support for model operations * Support is_causal as kwarg for SDPA * Update src/frontends/pytorch/src/op/addcmul.cpp * Update tests/model_hub_tests/torch_tests/test_timm.py * Support only decoder passed to convert_model * Fix tests * Update src/bindings/python/src/openvino/frontend/pytorch/fx_decoder.py * Update src/bindings/python/src/openvino/frontend/pytorch/fx_decoder.py Co-authored-by: Ekaterina Aidova * Apply suggestions from code review * Apply suggestions from code review * Improve testing with caching model * Apply suggestions from code review --------- Co-authored-by: Ekaterina Aidova --- .github/workflows/job_python_unit_tests.yml | 9 + .../openvino/frontend/pytorch/fx_decoder.py | 221 ++++++++++----- .../openvino/frontend/pytorch/ts_decoder.py | 12 + .../src/openvino/frontend/pytorch/utils.py | 5 + .../pyopenvino/frontend/pytorch/decoder.hpp | 17 +- .../python/src/pyopenvino/utils/utils.cpp | 2 + .../openvino/frontend/pytorch/decoder.hpp | 18 +- .../frontend/pytorch/node_context.hpp | 31 ++- src/frontends/pytorch/src/frontend.cpp | 1 + src/frontends/pytorch/src/node_context.cpp | 2 +- .../pytorch/src/op/adaptive_poolnd.cpp | 15 + src/frontends/pytorch/src/op/add.cpp | 8 +- src/frontends/pytorch/src/op/addcmul.cpp | 21 +- src/frontends/pytorch/src/op/addmm.cpp | 34 ++- src/frontends/pytorch/src/op/batch_norm.cpp | 132 +++++---- src/frontends/pytorch/src/op/div.cpp | 46 +++- src/frontends/pytorch/src/op/expand.cpp | 2 +- src/frontends/pytorch/src/op/full.cpp | 15 + src/frontends/pytorch/src/op/gelu.cpp | 29 +- src/frontends/pytorch/src/op/glu.cpp | 2 +- src/frontends/pytorch/src/op/index.cpp | 260 +++++++++++++++++- src/frontends/pytorch/src/op/leaky_relu.cpp | 32 +++ src/frontends/pytorch/src/op/mean.cpp | 26 +- src/frontends/pytorch/src/op/pad.cpp | 52 ++-- .../pytorch/src/op/scalar_tensor.cpp | 34 +++ .../src/op/scaled_dot_product_attention.cpp | 40 ++- .../pytorch/src/op/slice_scatter.cpp | 68 +++++ src/frontends/pytorch/src/op/split.cpp | 16 +- src/frontends/pytorch/src/op/sub.cpp | 16 +- src/frontends/pytorch/src/op/to.cpp | 10 + src/frontends/pytorch/src/op_table.cpp | 86 +++++- .../pytorch/src/translate_session.cpp | 19 +- src/frontends/pytorch/src/utils.cpp | 2 +- src/frontends/pytorch/src/utils.hpp | 15 +- tests/layer_tests/pytest.ini | 1 + .../pytorch_tests/pytorch_layer_test_class.py | 64 ++++- .../pytorch_tests/test_adaptive_avg_pool.py | 3 + .../pytorch_tests/test_adaptive_max_pool.py | 3 + .../layer_tests/pytorch_tests/test_addcmul.py | 1 + tests/layer_tests/pytorch_tests/test_addmm.py | 2 + tests/layer_tests/pytorch_tests/test_and.py | 4 + .../pytorch_tests/test_as_strided.py | 3 + .../pytorch_tests/test_batch_norm.py | 1 + .../pytorch_tests/test_broadcast_tensors.py | 1 + tests/layer_tests/pytorch_tests/test_clone.py | 1 + .../pytorch_tests/test_conv_transposend.py | 3 + .../pytorch_tests/test_convolution.py | 2 + .../layer_tests/pytorch_tests/test_device.py | 2 + tests/layer_tests/pytorch_tests/test_div.py | 2 + .../pytorch_tests/test_embedding.py | 1 + tests/layer_tests/pytorch_tests/test_eq.py | 1 + .../layer_tests/pytorch_tests/test_expand.py | 3 + .../layer_tests/pytorch_tests/test_flatten.py | 1 + .../pytorch_tests/test_floor_divide.py | 2 + tests/layer_tests/pytorch_tests/test_gelu.py | 1 + tests/layer_tests/pytorch_tests/test_glu.py | 1 + .../pytorch_tests/test_group_norm.py | 1 + tests/layer_tests/pytorch_tests/test_index.py | 27 ++ tests/layer_tests/pytorch_tests/test_mm.py | 3 + tests/layer_tests/pytorch_tests/test_mul.py | 2 + .../layer_tests/pytorch_tests/test_permute.py | 2 + .../layer_tests/pytorch_tests/test_reshape.py | 1 + .../pytorch_tests/test_reshape_as.py | 1 + .../pytorch_tests/test_resolve_conj_neg.py | 2 + .../test_scaled_dot_product_attention.py | 22 +- .../layer_tests/pytorch_tests/test_select.py | 1 + tests/layer_tests/pytorch_tests/test_silu.py | 1 + .../layer_tests/pytorch_tests/test_softmax.py | 2 + tests/layer_tests/pytorch_tests/test_sub.py | 1 + .../pytorch_tests/test_tensor_split.py | 1 + .../layer_tests/pytorch_tests/test_type_as.py | 1 + .../pytorch_tests/test_unflatten.py | 2 + .../layer_tests/pytorch_tests/test_unfold.py | 1 + .../torch_tests/requirements.txt | 1 + .../torch_tests/test_aliked.py | 2 +- .../torch_tests/test_detectron2.py | 2 +- .../model_hub_tests/torch_tests/test_edsr.py | 2 +- .../torch_tests/test_gfpgan.py | 2 +- .../torch_tests/test_hf_transformers.py | 2 +- .../torch_tests/test_speech-transformer.py | 2 +- .../model_hub_tests/torch_tests/test_timm.py | 10 +- .../torch_tests/test_torchbench.py | 2 +- .../torch_tests/test_torchvision_models.py | 13 +- .../torch_tests/torch_utils.py | 43 ++- tools/ovc/openvino/tools/ovc/convert_impl.py | 5 +- .../moc_frontend/pytorch_frontend_utils.py | 9 +- 86 files changed, 1283 insertions(+), 284 deletions(-) create mode 100644 src/frontends/pytorch/src/op/leaky_relu.cpp create mode 100644 src/frontends/pytorch/src/op/scalar_tensor.cpp create mode 100644 src/frontends/pytorch/src/op/slice_scatter.cpp diff --git a/.github/workflows/job_python_unit_tests.yml b/.github/workflows/job_python_unit_tests.yml index a04e2004eaf..e20a055eb5b 100644 --- a/.github/workflows/job_python_unit_tests.yml +++ b/.github/workflows/job_python_unit_tests.yml @@ -216,6 +216,15 @@ jobs: TEST_DEVICE: CPU TEST_PRECISION: FP32 + - name: PyTorch torch.export Layer Tests + if: ${{ fromJSON(inputs.affected-components).PyTorch_FE.test && runner.arch != 'ARM64' }} # Ticket: 126287 + run: | + python3 -m pytest ${LAYER_TESTS_INSTALL_DIR}/pytorch_tests -m precommit_torch_export --junitxml=${INSTALL_TEST_DIR}/TEST-pytorch.xml + env: + TEST_DEVICE: CPU + TEST_PRECISION: FP32 + PYTORCH_TRACING_MODE: EXPORT + - name: PyTorch torch.compile TORCHFX Layer Tests if: ${{ fromJSON(inputs.affected-components).PyTorch_FE.test && runner.os != 'macOS' }} run: | diff --git a/src/bindings/python/src/openvino/frontend/pytorch/fx_decoder.py b/src/bindings/python/src/openvino/frontend/pytorch/fx_decoder.py index 09375329cb7..cb6b6beb0e2 100644 --- a/src/bindings/python/src/openvino/frontend/pytorch/fx_decoder.py +++ b/src/bindings/python/src/openvino/frontend/pytorch/fx_decoder.py @@ -11,6 +11,7 @@ from openvino.frontend.pytorch.utils import maybe_convert_max_int, make_constant import torch + class TorchFXPythonDecoder (Decoder): def __init__(self, pt_module, fx_gm, nodes=None, mark_node_callback=None, input_shapes=[], input_types=[]): @@ -23,56 +24,96 @@ class TorchFXPythonDecoder (Decoder): self.input_types = input_types self.input_shapes = input_shapes + self._input_signature = [] + self._output_names = [] + if issubclass(type(pt_module), torch.fx.graph_module.GraphModule): + self._input_is_list = None self._nodes = list(pt_module.graph.nodes) self._inputs = [] self._outputs = [] for i in range(len(self._nodes)): if self._nodes[i].op == 'placeholder': self._inputs.append(i) + self._input_signature.append(self._nodes[i].name) elif self._nodes[i].op == 'output': # Instead of putting output index, refer to its target args = self._nodes[i].args if isinstance(args[0], tuple): args = args[0] - for output in args: - self._outputs.append(self._nodes.index(output)) + if isinstance(args[0], dict): + for name, output in args[0].items(): + self._outputs.append(self._nodes.index(output)) + self._output_names.append(name) + else: + for output in args: + self._outputs.append(self._nodes.index(output)) elif issubclass(type(pt_module), torch.fx.Node): - - self._nodes = nodes # passed from outer context + self._nodes = nodes # passed from outer context # FIXME: Quadratic complexity nodes*nodes considering the outer loop over all nodes for i in range(len(self._nodes)): if self._nodes[i] == pt_module: self._outputs = [i] - # code constant or None input as a tuple to diffirentiate it from regular input - # negative index is used to mark inputs initialized by inline constants, there are no such inputs in the graph - self._inputs = [self._nodes.index(arg) if arg in self._nodes else (arg,) for arg in pt_module.args] + # None in inputs mean the input is inlined or None (also considered inlined) + self._inputs = [self._nodes.index( + arg) if arg in self._nodes else (arg,) for arg in pt_module.args] - # FIXME: Find a better way to pass nested tuples to OV frontend. This is a temprary solution to flatten arguments. + # FIXME: Find a better way to pass nested tuples to OV frontend. This is a temporary solution to flatten arguments. new_inputs = [] for i in range(len(pt_module.args)): - expand_list = False - if isinstance(pt_module.args[i], list): + if isinstance(pt_module.args[i], list) and any([isinstance(a, torch.fx.Node) for a in pt_module.args[i]]): for arg in pt_module.args[i]: if arg in self._nodes: - expand_list = True - break; - if expand_list: - for arg in pt_module.args[i]: - new_inputs.append(self._nodes.index(arg)) + new_inputs.append(self._nodes.index(arg)) + else: + new_inputs.append((arg,)) else: new_inputs.append(self._inputs[i]) self._inputs = new_inputs - - def inputs(self): - return [x if x is not None else 100000 for x in self._inputs] + # Consider 0 a special case which may mean the input is inlined, but not guaranteed + return [x if not isinstance(x, tuple) else 0 for x in self._inputs] + + def is_input_inlined(self, index): + return isinstance(self._inputs[index], tuple) + + @staticmethod + def arg_to_constant(arg): + if isinstance(arg, list): + if len(arg) > 0: + return make_constant(pt_to_ov_type_map[type( + arg[0]).__name__], Shape([len(arg)]), arg) + else: + # TODO: which type should we use if list is empty? Need a signaling value here + return make_constant(int, Shape([0]), []) + elif isinstance(arg, bool): + return make_constant(OVType.boolean, Shape([]), [arg]) + elif isinstance(arg, int): + arg = maybe_convert_max_int(arg) + return make_constant(OVType.i32, Shape( + []), [arg]) # TODO: i32? why not i64? + elif isinstance(arg, float): + return make_constant(OVType.f32, Shape( + []), [arg]) # TODO: f32? why not f64? + return None + + + def inlined_input(self, index): + assert index < len(self._inputs), "Requested input doesn't exist" + assert isinstance(self._inputs[index], tuple), "Requested input which is not inlined" + assert self._inputs[index][0] is not None, "Requested None inlined input" + constant = None + arg = self._inputs[index][0] + constant = self.arg_to_constant(arg) + + assert constant is not None, "Constant wasn't created for inlined input" + return constant.outputs() def input(self, index): # TODO: remove return self.inputs()[index] # TODO: find specialized method @@ -81,6 +122,8 @@ class TorchFXPythonDecoder (Decoder): return "input"+str(index) def get_input_signature_name(self, index: int) -> str: + if self._input_signature is not None and index < len(self._input_signature): + return self._input_signature[index] return self.get_input_debug_name(index) def get_input_shape(self, index): @@ -89,6 +132,16 @@ class TorchFXPythonDecoder (Decoder): input = self._raw_input(index) return self.get_shape_for_value(input) + def get_input_strides(self, index: int) -> list: + raw_input = self._raw_input(index) + if isinstance(raw_input, torch.fx.node.Node) and hasattr(raw_input, "meta"): + meta = raw_input.meta + if "tensor_meta" in meta and hasattr(meta["tensor_meta"], "stride"): + strides = list(meta["tensor_meta"].stride) + if strides: + return strides + return [] + def get_input_type(self, index): if index < len(self.input_types): return OVAny(pt_to_ov_type_map[str(self.input_types[index])]) @@ -96,7 +149,10 @@ class TorchFXPythonDecoder (Decoder): return self.get_type_for_value(input) def get_output_debug_name(self, index): - return "output"+str(index) + if self._output_names is not None and index < len(self._output_names): + return self._output_names[index] + name = getattr(self.pt_module, "name", "output") + return name + ":" + str(index) def get_output_shape(self, index): output = self._raw_output(index) @@ -106,52 +162,58 @@ class TorchFXPythonDecoder (Decoder): output = self._raw_output(index) return self.get_type_for_value(output) - def _get_known_type_for_value(self, type): - ''' - Returns known/unknown types wrapped as OVAny - ''' - # Check for simple scalar types first - # TODO: Don't use str, use native types - if type is None: - return OVAny(OVType.dynamic) - if str(type) in pt_to_ov_type_map: - return OVAny(pt_to_ov_type_map[str(type)]) - elif type.__class__ is torch.TensorType: - # Tensor type, parse element type - # TODO: replace string by native type - # return OVAny(PartialShape([1,2,3])) - return OVAny(DecoderType.Tensor(self._get_known_type_for_value(type.dtype()))) - elif type.__class__ is torch.ListType: - element_type = type.getElementType() - return OVAny(DecoderType.List(self._get_known_type_for_value(element_type))) - else: - # Not yet recognized - return OVAny(OVType.dynamic) - #pt_type_class = value.type().__class__ - # if pt_type_class is torch.ListType: - def get_shape_for_value(self, value): - if value and ('tensor_meta' in value.meta.keys()): - return PartialShape(value.meta['tensor_meta'].shape) - return PartialShape([1]) + if value and hasattr(value, "meta") and ('tensor_meta' in value.meta.keys()): + if value.meta['tensor_meta']: + return PartialShape(len(value.meta['tensor_meta'].shape) * [-1]) + return PartialShape.dynamic() def get_type_for_value(self, value): if issubclass(type(value), torch.fx.Node): if ('tensor_meta' in value.meta.keys()): - pt_type = value.meta['tensor_meta'].dtype - if str(pt_type) in pt_to_ov_type_map: - ov_type = pt_to_ov_type_map[str(pt_type)] - return OVAny(ov_type) + if value.meta['tensor_meta'] and isinstance(value.meta['tensor_meta'], torch.Tensor): + pt_type = value.meta['tensor_meta'].dtype + if str(pt_type) in pt_to_ov_type_map: + ov_type = pt_to_ov_type_map[str(pt_type)] + return OVAny(ov_type) else: - return OVAny(OVType.f32) + return OVAny(OVType.dynamic) elif isinstance(value, int): return OVAny(OVType.i32) elif isinstance(value, float): return OVAny(OVType.f32) elif isinstance(value, bool): return OVAny(OVType.boolean) - else: - return OVAny(OVType.f32) + return OVAny(OVType.dynamic) + + def get_attribute(self, name): + if name in self.pt_module.kwargs: + attr = self.pt_module.kwargs[name] + if isinstance(attr, torch.dtype): + return OVAny(pt_to_ov_type_map[str(attr)]) + if isinstance(attr, torch.device): + return OVAny(attr.type) + if isinstance(attr, str): + return OVAny(attr) + # Numeric attrs convert to Constant + constant = self.arg_to_constant(attr) + if constant is not None: + return OVAny(constant.output(0)) + # so that has_attribute return True if attribute exist + return OVAny(DecoderType.PyNone()) + return OVAny(None) + + def get_named_input(self, name): + """ + Returns id of kwargs input. Such input can be Node or a constant value, + this function is only used for to return node index. If the input is + constant, get_attribute should be used. + """ + if name in self.pt_module.kwargs: + arg = self.pt_module.kwargs[name] + if isinstance(arg, torch.fx.Node): + return self._nodes.index(arg) + raise RuntimeError("This input is not a Node") def get_subgraph_size(self): if issubclass(type(self.pt_module), torch.fx.Node): @@ -165,8 +227,11 @@ class TorchFXPythonDecoder (Decoder): # make sure topological order is satisfied for node in self._nodes: if node.op == 'placeholder' or node.op == 'output': - continue # skipping non-operational nodes - decoder = TorchFXPythonDecoder(node, self.fx_gm, self._nodes, mark_node_callback=self.mark_node_callback) + continue # skipping non-operational nodes + if node.op == 'call_function' and str(node.target) in ["aten._assert_async.msg"]: + continue + decoder = TorchFXPythonDecoder( + node, self.fx_gm, self._nodes, mark_node_callback=self.mark_node_callback) self.m_decoders.append(decoder) node_visitor(decoder) @@ -176,7 +241,8 @@ class TorchFXPythonDecoder (Decoder): return list(self.pt_module.blocks()) def get_subgraph_decoder(self, index): - decoder = TorchFXPythonDecoder(self.get_subgraphs()[index], self.fx_gm, mark_node_callback=self.mark_node_callback) + decoder = TorchFXPythonDecoder(self.get_subgraphs( + )[index], self.fx_gm, mark_node_callback=self.mark_node_callback) self.m_decoders.append(decoder) return decoder @@ -202,7 +268,7 @@ class TorchFXPythonDecoder (Decoder): return self._raw_outputs()[index] def _raw_inputs(self): - return [self._nodes[x] if x is not None and x < len(self._nodes) else x for x in self._inputs] + return [self._nodes[x] if not isinstance(x, tuple) and x < len(self._nodes) else x[0] for x in self._inputs] def _raw_input(self, index): return self._raw_inputs()[index] @@ -214,6 +280,10 @@ class TorchFXPythonDecoder (Decoder): return self.outputs()[index] def mark_node(self, node): + name = self.get_op_type() + if "FrameworkNode" not in node.get_type_name(): + name += "/" + node.get_type_name() + node.set_friendly_name(name) if self.mark_node_callback is not None: self.mark_node_callback(self, node) return node @@ -223,10 +293,9 @@ class TorchFXPythonDecoder (Decoder): if self.pt_module.op == 'get_attr': # Extract Constant from FX module field ret = fetch_attr(self.fx_gm, self.pt_module.target) - ov_const = op.Constant(ret.numpy(), shared_memory=True) + ov_const = op.Constant(ret.numpy(force=True), shared_memory=True) return ov_const.outputs() - if not self.get_op_type() == 'prim::Constant': return None pt_value = self._raw_output(0) @@ -258,7 +327,8 @@ class TorchFXPythonDecoder (Decoder): ivalue = pt_value.toIValue() if pt_value.isCompleteTensor(): try: - ivalue = ivalue.to(memory_format=torch.contiguous_format).detach().cpu() + ivalue = ivalue.to( + memory_format=torch.contiguous_format).detach().cpu() except: print("[ WARNING ] Tensor couldn't detach") if str(pt_value.type().dtype()) in pt_to_py_type_map: @@ -273,12 +343,15 @@ class TorchFXPythonDecoder (Decoder): # this is only possible with adding a new ctor for Constant Python binding # TODO Check strides and pass them somehow values = ivalue.data_ptr() - ov_const = make_constant(ovtype, ovshape.get_shape(), values) + ov_const = make_constant( + ovtype, ovshape.get_shape(), values) except: # old variant that makes a slow data copying - print(f"[ WARNING ] Constant wasn't able to convert from data_ptr.") + print( + f"[ WARNING ] Constant wasn't able to convert from data_ptr.") values = ivalue.flatten().tolist() - ov_const = make_constant(ovtype, ovshape.get_shape(), values) + ov_const = make_constant( + ovtype, ovshape.get_shape(), values) return ov_const.outputs() else: # Incomplete tensor can be scalar @@ -294,14 +367,17 @@ class TorchFXPythonDecoder (Decoder): try: ovshape = PartialShape(ivalue.size()) ovtype = pt_to_ov_type_map[str(ivalue.type())] - ov_const = make_constant(ovtype, ovshape.get_shape(), ivalue.data_ptr()) + ov_const = make_constant( + ovtype, ovshape.get_shape(), ivalue.data_ptr()) except: # old variant that makes a slow data copying - print(f"[ WARNING ] Constant wasn't able to convert from data_ptr.") - nvalues = ivalue.numpy() + print( + f"[ WARNING ] Constant wasn't able to convert from data_ptr.") + nvalues = ivalue.numpy(force=True) ovtype = np_to_ov_type_map[str(nvalues.dtype)] ovshape = PartialShape(nvalues.shape) - ov_const = make_constant(ovtype, ovshape.get_shape(), nvalues.flatten().tolist()) + ov_const = make_constant( + ovtype, ovshape.get_shape(), nvalues.flatten().tolist()) return ov_const.outputs() return None @@ -325,7 +401,7 @@ class TorchFXPythonDecoder (Decoder): return ov_const.outputs() def input_is_none(self, index): - if index >= len(self.inputs()) or self._raw_input(index) is None: + if index >= len(self._inputs) or (isinstance(self._inputs[index], tuple) and self._inputs[index][0] is None): return True else: r_input = self._raw_input(index) @@ -342,7 +418,8 @@ class TorchFXPythonDecoder (Decoder): arg = self._inputs[i][0] if isinstance(arg, list): if len(arg) > 0: - constant = make_constant(pt_to_ov_type_map[type(arg[0]).__name__], Shape([len(arg)]), arg) + constant = make_constant(pt_to_ov_type_map[type( + arg[0]).__name__], Shape([len(arg)]), arg) else: # TODO: which type should we use if list is empty? Need a signaling value here constant = make_constant(int, Shape([0]), []) @@ -350,9 +427,11 @@ class TorchFXPythonDecoder (Decoder): constant = make_constant(OVType.boolean, Shape([]), [arg]) elif isinstance(arg, int): arg = maybe_convert_max_int(arg) - constant = make_constant(OVType.i32, Shape([]), [arg]) # TODO: i32? why not i64? + constant = make_constant(OVType.i32, Shape( + []), [arg]) # TODO: i32? why not i64? elif isinstance(arg, float): - constant = make_constant(OVType.f32, Shape([]), [arg]) # TODO: f32? why not f64? + constant = make_constant(OVType.f32, Shape( + []), [arg]) # TODO: f32? why not f64? if constant is None: if arg is None: diff --git a/src/bindings/python/src/openvino/frontend/pytorch/ts_decoder.py b/src/bindings/python/src/openvino/frontend/pytorch/ts_decoder.py index 7c2c53e68b0..bad0ce8633f 100644 --- a/src/bindings/python/src/openvino/frontend/pytorch/ts_decoder.py +++ b/src/bindings/python/src/openvino/frontend/pytorch/ts_decoder.py @@ -406,6 +406,18 @@ class TorchScriptPythonDecoder (Decoder): def inlined_inputs(self, index): return [] + def inlined_input(self, index): + return [] + + def is_input_inlined(self, index): + return False + + def get_attribute(self, name): + return OVAny(None) + + def get_named_input(self, name): + raise RuntimeError("There is no named inputs in TS graph") + @staticmethod def _transform_tensor_list_constants_to_listconstruct(graph: torch.Graph): # Function replaces prim::Constant containing List of Tensors with diff --git a/src/bindings/python/src/openvino/frontend/pytorch/utils.py b/src/bindings/python/src/openvino/frontend/pytorch/utils.py index 0225603e78a..157592f3aab 100644 --- a/src/bindings/python/src/openvino/frontend/pytorch/utils.py +++ b/src/bindings/python/src/openvino/frontend/pytorch/utils.py @@ -137,8 +137,13 @@ pt_to_ov_type_map = { "torch.bool": OVType.boolean, "torch.DoubleTensor": OVType.f64, "torch.FloatTensor": OVType.f32, + "torch.HalfTensor": OVType.f16, + "torch.BFloat16Tensor": OVType.bf16, "torch.IntTensor": OVType.i32, "torch.LongTensor": OVType.i64, + "torch.ShortTensor": OVType.i16, + "torch.CharTensor": OVType.i8, + "torch.ByteTensor": OVType.u8, "torch.BoolTensor": OVType.boolean, "torch.quint8": OVType.u8, "torch.qint8": OVType.i8, diff --git a/src/bindings/python/src/pyopenvino/frontend/pytorch/decoder.hpp b/src/bindings/python/src/pyopenvino/frontend/pytorch/decoder.hpp index 9e372203158..91c3c979d79 100644 --- a/src/bindings/python/src/pyopenvino/frontend/pytorch/decoder.hpp +++ b/src/bindings/python/src/pyopenvino/frontend/pytorch/decoder.hpp @@ -110,8 +110,21 @@ class PyDecoder : public ov::frontend::pytorch::TorchDecoder { PYBIND11_OVERRIDE_PURE(bool, TorchDecoder, may_produce_alias, in_index, out_index); } - ov::OutputVector inlined_inputs(size_t start_index) const override { - PYBIND11_OVERRIDE_PURE(ov::OutputVector, TorchDecoder, inlined_inputs, start_index); } + ov::OutputVector inlined_input(size_t index) const override { + PYBIND11_OVERRIDE_PURE(ov::OutputVector, TorchDecoder, inlined_input, index); + } + + bool is_input_inlined(size_t index) const override { + PYBIND11_OVERRIDE_PURE(bool, TorchDecoder, is_input_inlined, index); + } + + ov::Any get_attribute(const std::string &name) const override{ + PYBIND11_OVERRIDE_PURE(ov::Any, TorchDecoder, get_attribute, name); + } + + size_t get_named_input(const std::string &name) const override{ + PYBIND11_OVERRIDE_PURE(size_t, TorchDecoder, get_named_input, name); + } const std::string& decoder_type_name() const override { PYBIND11_OVERRIDE_PURE(const std::string&, TorchDecoder, decoder_type_name); diff --git a/src/bindings/python/src/pyopenvino/utils/utils.cpp b/src/bindings/python/src/pyopenvino/utils/utils.cpp index 906c19ae820..57a17ad6fb5 100644 --- a/src/bindings/python/src/pyopenvino/utils/utils.cpp +++ b/src/bindings/python/src/pyopenvino/utils/utils.cpp @@ -375,6 +375,8 @@ ov::Any py_object_to_any(const py::object& py_obj) { return py::cast(py_obj); } else if (py::isinstance(py_obj)) { return py::cast(py_obj); + } else if (py::isinstance>(py_obj)) { + return py::cast>(py_obj); // FrontEnd Decoder } else if (py::isinstance(py_obj)) { return py::cast>(py_obj); diff --git a/src/frontends/pytorch/include/openvino/frontend/pytorch/decoder.hpp b/src/frontends/pytorch/include/openvino/frontend/pytorch/decoder.hpp index d5878783c31..43243ab6c64 100644 --- a/src/frontends/pytorch/include/openvino/frontend/pytorch/decoder.hpp +++ b/src/frontends/pytorch/include/openvino/frontend/pytorch/decoder.hpp @@ -19,7 +19,7 @@ public: // fundamental types like int, float etc. virtual Any const_input(size_t index) const = 0; - // Using size_t for input/output unuque ids are in sync with torch code, see def in + // Using size_t for input/output unique ids are in sync with torch code, see def in // torch/include/torch/csrc/jit/ir/ir.h, Value::unique_ // TODO: set of input and output methods are not aligned; also they are not aligned with the rest of FEs @@ -89,7 +89,7 @@ public: virtual size_t output(size_t index) const = 0; // Embed mapping to/from the original node representation from/to node passed as a parameter - // the representation of this mapping is specific for particular decored type and may be NOP + // the representation of this mapping is specific for particular decorated type and may be NOP // returns the same node as syntactically convenient way to make nested sentences in code virtual std::shared_ptr mark_node(std::shared_ptr ov_node) const = 0; @@ -109,9 +109,19 @@ public: /// Returns new nodes for inputs inlined in the op itself // Used in Torch.FX decoder - virtual OutputVector inlined_inputs(size_t start_index) const = 0; + virtual OutputVector inlined_input(size_t index) const = 0; - /// Returns the id of the deccoder type (0: TorchFX, 1: TorchScript) + /// Returns if input is inlined + // Used in Torch.FX decoder + virtual bool is_input_inlined(size_t index) const = 0; + + /// Returns named attribute as Any. For example kwargs input for FX graph + virtual ov::Any get_attribute(const std::string& name) const = 0; + + /// Returns index of named input. For example kwargs input for FX graph + virtual size_t get_named_input(const std::string& name) const = 0; + + /// Returns the id of the decoder type ("fx": TorchFX, "ts": TorchScript) virtual const std::string& decoder_type_name() const = 0; }; diff --git a/src/frontends/pytorch/include/openvino/frontend/pytorch/node_context.hpp b/src/frontends/pytorch/include/openvino/frontend/pytorch/node_context.hpp index ae8f0b34a86..307f04391d4 100644 --- a/src/frontends/pytorch/include/openvino/frontend/pytorch/node_context.hpp +++ b/src/frontends/pytorch/include/openvino/frontend/pytorch/node_context.hpp @@ -51,12 +51,36 @@ public: // Search for input in tensor map and return an output port for already converted op // TODO: int due to base class uses it, but naturally it should be size_t for PT Output get_input(int index) const override { - FRONT_END_GENERAL_CHECK(!input_is_none(index), "Input is none with index: ", index); + size_t index_ = static_cast(index); + FRONT_END_GENERAL_CHECK(!input_is_none(index_), "Input doesn't exist with index: ", index); auto input = m_decoder_inputs.at(index); + if (input == 0) { + // Case when input can be inlined (possible only for fx decoder) + if (m_decoder->is_input_inlined(index_)) { + auto inlined_input = m_decoder->inlined_input(index_); + FRONT_END_GENERAL_CHECK(inlined_input.size() == 1, "Incorrect inlined input with index:", index); + return inlined_input[0]; + } + } FRONT_END_GENERAL_CHECK(m_tensor_map->count(input), "No tensor corresponding input: ", input, " exist."); return m_tensor_map->at(input); } + Output get_input(const std::string& name) const override { + FRONT_END_GENERAL_CHECK(has_attribute(name), "Input with name ", name, " doesn't exist"); + auto attr = get_attribute_as_any(name); + if (attr.is>()) { + // Case when input is constant value + return attr.as>(); + } else if (attr.is()) { + // None means input is unknown type, most likely a Node + auto input = m_decoder->get_named_input(name); + FRONT_END_GENERAL_CHECK(m_tensor_map->count(input), "No tensor corresponding input: ", input, " exist."); + return m_tensor_map->at(input); + } + FRONT_END_GENERAL_CHECK(false, "Input has type which can't be converted to ov::Node."); + } + Any get_values_from_const_input(int index) const override; // TODO: upstream to base class @@ -112,9 +136,8 @@ public: return ov_output; } - Any get_attribute_as_any(const std::string&) const override { - throw std::runtime_error( - "There is no any named attributes in PyTorch node, query by attribute name is not implemented"); + Any get_attribute_as_any(const std::string& name) const override { + return m_decoder->get_attribute(name); } void mutate_input(size_t index, Output ov_output) const; diff --git a/src/frontends/pytorch/src/frontend.cpp b/src/frontends/pytorch/src/frontend.cpp index ee9a9919045..fa00b94a4c6 100644 --- a/src/frontends/pytorch/src/frontend.cpp +++ b/src/frontends/pytorch/src/frontend.cpp @@ -194,6 +194,7 @@ void FrontEnd::normalize(const std::shared_ptr& model) const { manager.register_pass(); manager.register_pass(); manager.register_pass(); + // TODO: remove AtenIndexToSelect when problem with dynamic input rank is gone. manager.register_pass(); manager.register_pass(); manager.register_pass(); diff --git a/src/frontends/pytorch/src/node_context.cpp b/src/frontends/pytorch/src/node_context.cpp index 0ebd26d3b8e..ff6bb045e7f 100644 --- a/src/frontends/pytorch/src/node_context.cpp +++ b/src/frontends/pytorch/src/node_context.cpp @@ -109,7 +109,7 @@ Output NodeContext::get_tensor_from_model_or_create_input(size_t index) co } Output NodeContext::get_input_from_visible_context(size_t index) const { - FRONT_END_GENERAL_CHECK(index < get_input_size(), "Index is lower then number of inputs."); + FRONT_END_GENERAL_CHECK(index < get_input_size(), "Index ", index, " is lower then number of inputs."); auto input_tensor = get_input(static_cast(index)); auto input_node = input_tensor.get_node_shared_ptr(); if (std::dynamic_pointer_cast(input_node)) { diff --git a/src/frontends/pytorch/src/op/adaptive_poolnd.cpp b/src/frontends/pytorch/src/op/adaptive_poolnd.cpp index 9c349e50c8e..bc6b17cbed0 100644 --- a/src/frontends/pytorch/src/op/adaptive_poolnd.cpp +++ b/src/frontends/pytorch/src/op/adaptive_poolnd.cpp @@ -117,6 +117,21 @@ OutputVector translate_adaptive_max_pool1d(const NodeContext& context) { return translate_adaptive_max_pool_base(context, const_tile_params, const_neg_1); }; +OutputVector translate_adaptive_max_pool3d_fx(const NodeContext& context) { + auto outs = translate_adaptive_max_pool3d(context); + return {context.mark_node(make_list_construct(outs))}; +}; + +OutputVector translate_adaptive_max_pool2d_fx(const NodeContext& context) { + auto outs = translate_adaptive_max_pool2d(context); + return {context.mark_node(make_list_construct(outs))}; +}; + +OutputVector translate_adaptive_max_pool1d_fx(const NodeContext& context) { + auto outs = translate_adaptive_max_pool1d(context); + return {context.mark_node(make_list_construct(outs))}; +}; + } // namespace op } // namespace pytorch } // namespace frontend diff --git a/src/frontends/pytorch/src/op/add.cpp b/src/frontends/pytorch/src/op/add.cpp index 33699ad90fa..ab2933bd6f7 100644 --- a/src/frontends/pytorch/src/op/add.cpp +++ b/src/frontends/pytorch/src/op/add.cpp @@ -34,8 +34,14 @@ OutputVector translate_add_common(const NodeContext& context, bool inplace) { } else { align_eltwise_input_types(context, lhs, rhs, true); } + Output alpha; if (!context.input_is_none(2)) { - auto converted_alpha = context.mark_node(std::make_shared(context.get_input(2), rhs)); + alpha = context.get_input(2); + } else if (context.has_attribute("alpha")) { + alpha = context.get_attribute>("alpha"); + } + if (alpha.get_node_shared_ptr()) { + auto converted_alpha = context.mark_node(std::make_shared(alpha, rhs)); rhs = context.mark_node(std::make_shared(converted_alpha, rhs)); } auto add = context.mark_node(std::make_shared(lhs, rhs)); diff --git a/src/frontends/pytorch/src/op/addcmul.cpp b/src/frontends/pytorch/src/op/addcmul.cpp index 7933acfde1d..9ebe2a24875 100644 --- a/src/frontends/pytorch/src/op/addcmul.cpp +++ b/src/frontends/pytorch/src/op/addcmul.cpp @@ -17,15 +17,30 @@ namespace op { using namespace ov::op; -OutputVector translate_addcmul(const NodeContext& context) { - num_inputs_check(context, 4, 4); +namespace { +OutputVector addcmul_common(const NodeContext& context, const Output& value) { const auto eltwise_mult = std::make_shared(context.get_input(1), context.get_input(2)); - const auto value = context.get_input(3); const auto converted_value = std::make_shared(value, context.get_input(1)); const auto scalar_mult = std::make_shared(eltwise_mult, converted_value); context.mark_nodes({eltwise_mult, converted_value, scalar_mult}); return {context.mark_node(std::make_shared(context.get_input(0), scalar_mult))}; }; +} // namespace + +OutputVector translate_addcmul(const NodeContext& context) { + num_inputs_check(context, 4, 4); + const auto value = context.get_input(3); + return addcmul_common(context, value); +}; + +OutputVector translate_addcmul_fx(const NodeContext& context) { + num_inputs_check(context, 3, 3); + Output value = context.mark_node(v0::Constant::create(element::f32, Shape{}, {1})); + if (context.has_attribute("value")) { + value = context.get_input("value"); + } + return addcmul_common(context, value); +}; } // namespace op } // namespace pytorch diff --git a/src/frontends/pytorch/src/op/addmm.cpp b/src/frontends/pytorch/src/op/addmm.cpp index 731014c6c27..097737ec93e 100644 --- a/src/frontends/pytorch/src/op/addmm.cpp +++ b/src/frontends/pytorch/src/op/addmm.cpp @@ -16,12 +16,22 @@ namespace op { using namespace ov::op; -OutputVector translate_addmm(const NodeContext& context) { - num_inputs_check(context, 3, 5); +namespace { +OutputVector translate_addmm_common(const NodeContext& context, const Output beta, const Output alpha) { auto input = context.get_input(0); auto m1 = context.get_input(1); auto m2 = context.get_input(2); auto mm = context.mark_node(std::make_shared(m1, m2)); + auto beta_converted = context.mark_node(std::make_shared(beta, input)); + auto alpha_converted = context.mark_node(std::make_shared(alpha, mm)); + auto input_beta = context.mark_node(std::make_shared(input, beta_converted)); + auto mm_alpha = context.mark_node(std::make_shared(mm, alpha_converted)); + return {context.mark_node(std::make_shared(input_beta, mm_alpha))}; +}; +} // namespace + +OutputVector translate_addmm(const NodeContext& context) { + num_inputs_check(context, 3, 5); auto one = context.mark_node(v0::Constant::create(element::f32, Shape{}, {1})); ov::Output alpha = one; ov::Output beta = one; @@ -31,11 +41,21 @@ OutputVector translate_addmm(const NodeContext& context) { if (!context.input_is_none(4)) { alpha = context.get_input(4); } - auto beta_converted = context.mark_node(std::make_shared(beta, input)); - auto alpha_converted = context.mark_node(std::make_shared(alpha, mm)); - auto input_beta = context.mark_node(std::make_shared(input, beta_converted)); - auto mm_alpha = context.mark_node(std::make_shared(mm, alpha_converted)); - return {context.mark_node(std::make_shared(input_beta, mm_alpha))}; + return {translate_addmm_common(context, beta, alpha)}; +}; + +OutputVector translate_addmm_fx(const NodeContext& context) { + num_inputs_check(context, 3, 3); + auto one = context.mark_node(v0::Constant::create(element::f32, Shape{}, {1})); + ov::Output alpha = one; + ov::Output beta = one; + if (context.has_attribute("beta")) { + beta = context.get_input("beta"); + } + if (context.has_attribute("alpha")) { + alpha = context.get_input("alpha"); + } + return {translate_addmm_common(context, beta, alpha)}; }; } // namespace op diff --git a/src/frontends/pytorch/src/op/batch_norm.cpp b/src/frontends/pytorch/src/op/batch_norm.cpp index 126588eb952..5d10222eceb 100644 --- a/src/frontends/pytorch/src/op/batch_norm.cpp +++ b/src/frontends/pytorch/src/op/batch_norm.cpp @@ -17,6 +17,7 @@ #include "openvino/op/subtract.hpp" #include "openvino/op/unsqueeze.hpp" #include "openvino/op/util/framework_node.hpp" +#include "openvino/pass/graph_rewrite.hpp" #include "utils.hpp" namespace ov { @@ -37,91 +38,112 @@ Output broadcast_const_to_channel_dim(const NodeContext& context, auto channel_dim_exp = context.mark_node(std::make_shared(channel_dim, zero_i)); return context.mark_node(std::make_shared(value, channel_dim_exp)); } + +OutputVector make_batch_norm(const NodeContext& context, + const Output& input, + const Output& weight, + const Output& bias, + const Output& running_mean, + const Output& running_var, + float epsilon) { + Output w = weight; + Output b = bias; + Output mean = running_mean; + Output var = running_var; + if (!w.get_node_shared_ptr()) { + auto one_f = context.mark_node(v0::Constant::create(element::f32, Shape{}, {1})); + w = broadcast_const_to_channel_dim(context, input, one_f); + } + if (!b.get_node_shared_ptr()) { + auto zero_f = context.mark_node(v0::Constant::create(element::f32, Shape{}, {0})); + b = broadcast_const_to_channel_dim(context, input, zero_f); + } + auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); + auto zero_1d = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0})); + auto one = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1})); + auto two = context.mark_node(v0::Constant::create(element::i32, Shape{}, {2})); + Output rank = std::get<1>(get_shape_rank(context, input, true)); + auto after_channel_dims = context.mark_node(std::make_shared(two, rank, one)); + auto axes = context.mark_node(std::make_shared(OutputVector{zero_1d, after_channel_dims}, 0)); + if (!mean.get_node_shared_ptr()) { + mean = context.mark_node(std::make_shared(input, axes, false)); + } + if (!var.get_node_shared_ptr()) { + auto current_mean = context.mark_node(std::make_shared(input, axes, true)); + auto sub_v = context.mark_node(std::make_shared(input, current_mean)); + auto sqr_sub = context.mark_node(std::make_shared(sub_v, sub_v)); + var = context.mark_node(std::make_shared(sqr_sub, axes, false)); + } + return {context.mark_node(std::make_shared(input, w, b, mean, var, epsilon))}; +} } // namespace -OutputVector translate_batch_norm_common(const NodeContext& context, bool training) { +OutputVector translate_batch_norm(const NodeContext& context) { // Schema: aten::batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, // bool training, float momentum, float eps, bool cudnn_enabled) -> Tensor - - // batch_norm_legit_no_training Schema: aten::batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? - // running_mean, Tensor? running_var, float momentum, float eps) -> Tensor - - auto input = context.get_input(0); + num_inputs_check(context, 7, 9); Output weight; Output bias; Output running_mean; Output running_var; - Output current_mean; - Output current_var; if (!context.input_is_none(1)) { weight = context.get_input(1); - } else { - auto one_f = context.mark_node(v0::Constant::create(element::f32, Shape{}, {1})); - weight = broadcast_const_to_channel_dim(context, input, one_f); } if (!context.input_is_none(2)) { bias = context.get_input(2); - } else { - auto zero_f = context.mark_node(v0::Constant::create(element::f32, Shape{}, {0})); - bias = broadcast_const_to_channel_dim(context, input, zero_f); } // index 3 running_mean and index 4 running_var can be none for training case only, check that not training before // if training for batch norm activated, but model in eval mode, it uses current statistics instead of running - if (training) { - auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); - auto zero_1d = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0})); - auto one = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1})); - auto two = context.mark_node(v0::Constant::create(element::i32, Shape{}, {2})); - auto input_shape = context.mark_node(std::make_shared(input, element::i32)); - auto rank_unsq = context.mark_node(std::make_shared(input_shape, element::i32)); - auto rank = context.mark_node(std::make_shared(rank_unsq, zero)); - auto after_channel_dims = context.mark_node(std::make_shared(two, rank, one)); - auto axes = context.mark_node(std::make_shared(OutputVector{zero_1d, after_channel_dims}, 0)); - current_mean = context.mark_node(std::make_shared(input, axes, false)); - auto mean = context.mark_node(std::make_shared(input, axes, true)); - auto sub_v = context.mark_node(std::make_shared(input, mean)); - auto sqr_sub = context.mark_node(std::make_shared(sub_v, sub_v)); - current_var = context.mark_node(std::make_shared(sqr_sub, axes, false)); - } + auto training = context.const_input(5); if (!training) { running_mean = context.get_input(3); - } else { - running_mean = current_mean; - } - if (!training) { running_var = context.get_input(4); - } else { - running_var = current_var; } // Input with index 6 is momentum, it is used only for updating running_mean accumulation during training - // In batch_norm_legit_no_training, momentum is index 5 and epsilon is 6 - float epsilon; - if (context.get_input_size() == 7) { - epsilon = context.const_input(6); - } else { - epsilon = context.const_input(7); - } + float epsilon = context.const_input(7); // Input with index 8 is flag "cudnn_enabled" we can ignore it - return {context.mark_node( - std::make_shared(input, weight, bias, running_mean, running_var, epsilon))}; -}; - -OutputVector translate_batch_norm(const NodeContext& context) { - num_inputs_check(context, 7, 9); - auto training = context.const_input(5); - return translate_batch_norm_common(context, training); + return make_batch_norm(context, context.get_input(0), weight, bias, running_mean, running_var, epsilon); } OutputVector translate_batch_norm_legit_fx(const NodeContext& context) { - num_inputs_check(context, 7, 9); - auto training = context.const_input(5); - auto output = translate_batch_norm_common(context, training); + auto output = translate_batch_norm(context); return {context.mark_node(make_list_construct(output))}; } OutputVector translate_batch_norm_legit_no_training_fx(const NodeContext& context) { num_inputs_check(context, 7, 9); - auto output = translate_batch_norm_common(context, false); + Output weight; + Output bias; + if (!context.input_is_none(1)) { + weight = context.get_input(1); + } + if (!context.input_is_none(2)) { + bias = context.get_input(2); + } + auto running_mean = context.get_input(3); + auto running_var = context.get_input(4); + float epsilon = context.const_input(6); + auto output = make_batch_norm(context, context.get_input(0), weight, bias, running_mean, running_var, epsilon); + return {context.mark_node(make_list_construct(output))}; +} + +OutputVector translate_batch_norm_legit_no_stats_fx(const NodeContext& context) { + num_inputs_check(context, 6, 6); + // torch.ops.aten._native_batch_norm_legit.no_stats(arg2_1, arg0_1, arg1_1, True, 0.1, 5e-05) + Output weight; + if (!context.input_is_none(1)) { + weight = context.get_input(1); + } + Output bias; + if (!context.input_is_none(2)) { + bias = context.get_input(2); + } + auto training = context.const_input(3); + FRONT_END_OP_CONVERSION_CHECK(training, + "aten._native_batch_norm_legit.no_stats can only be used when training=True."); + // index 4 momentum is used during training only + auto eps = context.const_input(5); + auto output = make_batch_norm(context, context.get_input(0), weight, bias, {}, {}, eps); return {context.mark_node(make_list_construct(output))}; } diff --git a/src/frontends/pytorch/src/op/div.cpp b/src/frontends/pytorch/src/op/div.cpp index dbbb6c89af7..477c9ced2bb 100644 --- a/src/frontends/pytorch/src/op/div.cpp +++ b/src/frontends/pytorch/src/op/div.cpp @@ -17,14 +17,13 @@ namespace frontend { namespace pytorch { namespace op { -OutputVector translate_div_common(const NodeContext& context, bool inplace) { - num_inputs_check(context, 2, 3); - auto x = context.get_input(0); - auto y = context.get_input(1); - std::string rounding_mode = ""; - if (!context.input_is_none(2)) { - rounding_mode = context.const_input(2); - } +OutputVector translate_div_common(const NodeContext& context, + const Output& lhs, + const Output& rhs, + const std::string& rounding_mode, + bool inplace) { + auto x = lhs; + auto y = rhs; if (rounding_mode.empty()) { // if no rounding mode and both inputs are ints cast BOTH to fp32 const auto x_dtype = x.get_element_type(); @@ -36,7 +35,7 @@ OutputVector translate_div_common(const NodeContext& context, bool inplace) { } if (inplace) { if (x.get_element_type().is_dynamic() || x.get_element_type() != y.get_element_type()) - y = context.mark_node(std::make_shared(x, y)); + y = context.mark_node(std::make_shared(y, x)); } else { align_eltwise_input_types(context, x, y, true); } @@ -55,11 +54,36 @@ OutputVector translate_div_common(const NodeContext& context, bool inplace) { }; OutputVector translate_div(const NodeContext& context) { - return translate_div_common(context, false); + num_inputs_check(context, 2, 3); + auto x = context.get_input(0); + auto y = context.get_input(1); + std::string rounding_mode = ""; + if (!context.input_is_none(2)) { + rounding_mode = context.const_input(2); + } + return translate_div_common(context, x, y, rounding_mode, false); }; OutputVector translate_div_(const NodeContext& context) { - return translate_div_common(context, true); + num_inputs_check(context, 2, 3); + auto x = context.get_input(0); + auto y = context.get_input(1); + std::string rounding_mode = ""; + if (!context.input_is_none(2)) { + rounding_mode = context.const_input(2); + } + return translate_div_common(context, x, y, rounding_mode, true); +}; + +OutputVector translate_div_fx(const NodeContext& context) { + num_inputs_check(context, 2, 2); + auto x = context.get_input(0); + auto y = context.get_input(1); + std::string rounding_mode = ""; + if (context.has_attribute("rounding_mode")) { + rounding_mode = context.get_attribute("rounding_mode"); + } + return translate_div_common(context, x, y, rounding_mode, false); }; } // namespace op diff --git a/src/frontends/pytorch/src/op/expand.cpp b/src/frontends/pytorch/src/op/expand.cpp index 0b069040412..7fcb7a898a4 100644 --- a/src/frontends/pytorch/src/op/expand.cpp +++ b/src/frontends/pytorch/src/op/expand.cpp @@ -49,7 +49,7 @@ OutputVector translate_expand_fx(const NodeContext& context) { // output shapes are same. This should be removed after a proper optimization is // implemented. auto sizes_const = context.const_input(1); - if (x.get_shape() == sizes_const) { + if (x.get_partial_shape().is_static() && x.get_shape() == sizes_const) { return {x}; } auto sizes = context.get_input(1); diff --git a/src/frontends/pytorch/src/op/full.cpp b/src/frontends/pytorch/src/op/full.cpp index a86b5bf3725..70ff0ba980a 100644 --- a/src/frontends/pytorch/src/op/full.cpp +++ b/src/frontends/pytorch/src/op/full.cpp @@ -71,6 +71,21 @@ OutputVector translate_full(const NodeContext& context) { return {base_translate_full_with_convert(context, sizes, value, dtype_id)}; }; +OutputVector translate_full_fx(const NodeContext& context) { + // aten.full.default([16, 16], 0, dtype = torch.float32, layout = torch.strided, device = device(type='cpu'), + // pin_memory = False) + num_inputs_check(context, 2, 2); + auto sizes = context.get_input(0); + auto value = context.get_input(1); + + auto filled_tensor = base_translate_full(context, sizes, value); + if (context.has_attribute("dtype")) { + auto dtype = context.get_attribute("dtype"); + filled_tensor = context.mark_node(std::make_shared(filled_tensor, dtype)); + } + return {filled_tensor}; +}; + OutputVector translate_full_like(const NodeContext& context) { num_inputs_check(context, 2, 7); auto input = context.get_input(0); diff --git a/src/frontends/pytorch/src/op/gelu.cpp b/src/frontends/pytorch/src/op/gelu.cpp index 4f7fc810596..64afb511b31 100644 --- a/src/frontends/pytorch/src/op/gelu.cpp +++ b/src/frontends/pytorch/src/op/gelu.cpp @@ -12,13 +12,9 @@ namespace frontend { namespace pytorch { namespace op { -OutputVector translate_gelu(const NodeContext& context) { - num_inputs_check(context, 1, 2); +namespace { +OutputVector translate_gelu_common(const NodeContext& context, const std::string& approximate) { auto x = context.get_input(0); - std::string approximate = "none"; - if (!context.input_is_none(1)) { - approximate = context.const_input(1); - } if (approximate == "none") { return {context.mark_node(std::make_shared(x, ov::op::GeluApproximationMode::ERF))}; } @@ -27,6 +23,27 @@ OutputVector translate_gelu(const NodeContext& context) { } FRONT_END_OP_CONVERSION_CHECK(false, "Unsupported approximate for Gelu: ", approximate); }; +} // namespace + +OutputVector translate_gelu(const NodeContext& context) { + num_inputs_check(context, 1, 2); + auto x = context.get_input(0); + std::string approximate = "none"; + if (!context.input_is_none(1)) { + approximate = context.const_input(1); + } + return translate_gelu_common(context, approximate); +}; + +OutputVector translate_gelu_fx(const NodeContext& context) { + num_inputs_check(context, 1, 1); + auto x = context.get_input(0); + std::string approximate = "none"; + if (context.has_attribute("approximate")) { + approximate = context.get_attribute("approximate"); + } + return translate_gelu_common(context, approximate); +}; } // namespace op } // namespace pytorch diff --git a/src/frontends/pytorch/src/op/glu.cpp b/src/frontends/pytorch/src/op/glu.cpp index dbe979fb1f2..779c81c4904 100644 --- a/src/frontends/pytorch/src/op/glu.cpp +++ b/src/frontends/pytorch/src/op/glu.cpp @@ -17,7 +17,7 @@ namespace op { using namespace ov::op; OutputVector translate_glu(const NodeContext& context) { - num_inputs_check(context, 2, 2); + num_inputs_check(context, 1, 2); auto x = context.get_input(0); auto dim = context.input_is_none(1) ? context.mark_node(v0::Constant::create(element::i32, Shape{}, {-1})) : context.get_input(1); diff --git a/src/frontends/pytorch/src/op/index.cpp b/src/frontends/pytorch/src/op/index.cpp index e739476619f..a9aaf62257d 100644 --- a/src/frontends/pytorch/src/op/index.cpp +++ b/src/frontends/pytorch/src/op/index.cpp @@ -3,9 +3,20 @@ // #include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/add.hpp" #include "openvino/op/concat.hpp" #include "openvino/op/convert.hpp" +#include "openvino/op/gather.hpp" #include "openvino/op/gather_nd.hpp" +#include "openvino/op/multiply.hpp" +#include "openvino/op/non_zero.hpp" +#include "openvino/op/reduce_prod.hpp" +#include "openvino/op/reshape.hpp" +#include "openvino/op/shape_of.hpp" +#include "openvino/op/slice.hpp" +#include "openvino/op/split.hpp" +#include "openvino/op/transpose.hpp" +#include "pt_framework_node.hpp" #include "utils.hpp" namespace ov { @@ -13,23 +24,252 @@ namespace frontend { namespace pytorch { namespace op { +using namespace ov::op; + +namespace { +Output flatten(ov::pass::NodeRegistry& rg, const Output& value, size_t axis) { + // First dimension of output tensor is the product of [d_0, ... d_{axis-1}] dimensions of + // input tensor. The last dimension is the product of the rest of input tensor dimensions: + // [d_{axis}, ..., d_n] + Output output_shape; + if (axis == 0) { + output_shape = v0::Constant::create(element::i32, Shape{2}, {1, -1}); + } else if (axis == 1) { + output_shape = v0::Constant::create(element::i32, Shape{2}, {0, -1}); + } else { + const auto value_shape = rg.make(value, element::i32); + const auto value_rank = rg.make(value_shape, element::i32); + const auto axis_node = v0::Constant::create(element::i32, Shape{1}, {axis}); + auto start = v0::Constant::create(element::i32, Shape{1}, {0}); + auto step = v0::Constant::create(element::i32, Shape{1}, {1}); + const auto first_part_dims = rg.make(value_shape, start, axis_node, step); + auto zero = v0::Constant::create(element::i32, {}, {0}); + auto first_part_dims_length = rg.make(first_part_dims, zero, true); + + auto remaining_part_length = v0::Constant::create(element::i32, {1}, {-1}); + + output_shape = rg.make(OutputVector{first_part_dims_length, remaining_part_length}, 0); + } + return rg.make(value, output_shape, true); +} + +OutputVector index_on_list(ov::pass::NodeRegistry& rg, + const Output& data, + std::deque> ids, + int64_t rank) { + // Multiple tensors as indices. Each tensor could either be + // 1. prim::Constant() + // representing ":" in python indexing. E.g. tensor[:, :] + // 2. prim::Constant[value=...] or tensor output + // representing advanced indexing. E.g. tensor[[0, 1], [2, 0]]. + // For more info on advanced indexing, + // check https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#advanced-indexing + + // Consider a general case of + // t: [x_1, y_1, y_2, ..., x_m, ..., y_n] + // where t is a tensor of rank m+n, {x_i} are axes where tensor index is provided, and {y_i} are axes for + // ":". Same results can be achieved through transposing t into + // t: [x_1, x_2, ..., x_m, y_1, y_2, ..., y_n] + // and use gather + // t: [x_1 * x_2 * ... * x_m, y_1 * y_2 * ... * y_n] + // tensor index = \sum_{i=1}^m (ind_i * \prod_{j=i+1}^m (x_j)) + // After gather, reshape and transpose back. + std::vector advanced_ids; + std::vector is_masked_bool; + OutputVector masked_indicies; + // for case when index is bool e.g. x[x>0], replace index with non_zero + for (size_t i = 0; i < ids.size(); i++) { + // skip dimensions where index is None + bool is_none = false; + if (!ids[i].get_node_shared_ptr()) { + is_none = true; + } + if (auto const_input = cast_fw_node(ids[i].get_node_shared_ptr(), "prim::Constant")) { + const auto& attrs = const_input->get_attrs(); + if (attrs.find("none_value") != attrs.end()) { + is_none = true; + } + } + if (is_none) { + masked_indicies.push_back(ids[i]); + is_masked_bool.push_back(false); + continue; + } + auto id_dtype = ids[i].get_element_type(); + if (id_dtype == element::boolean || id_dtype == element::u8) { + auto idx = rg.make(ids[i], element::u8); + auto nonzero = rg.make(idx, element::i32); + auto input_order = v0::Constant::create(element::i32, Shape{2}, {1, 0}); + auto masked_id = rg.make(nonzero, input_order); + masked_indicies.push_back(masked_id); + is_masked_bool.push_back(true); + } else { + masked_indicies.push_back(ids[i]); + is_masked_bool.push_back(false); + } + advanced_ids.push_back(i); + } + + // all indicies prim::Constant(None), return input as is + if (advanced_ids.size() == 0) { + return {data}; + } + // perform gather for single element case + if (advanced_ids.size() == 1) { + auto index = masked_indicies[advanced_ids[0]]; + if (is_masked_bool[advanced_ids[0]]) { + auto gather = rg.make(data, index); + return {gather}; + } + index = rg.make(index, element::i32); + auto dim = v0::Constant::create(element::i32, Shape{}, {advanced_ids[0]}); + auto gather = rg.make(data, index, dim); + return {gather}; + } + auto adv_idx_count = advanced_ids.size(); + auto input_shape = rg.make(data, element::i32); + auto zero = v0::Constant::create(element::i32, Shape{}, {0}); + auto input_dims = rg.make(input_shape, zero, rank); + std::vector non_used_dims; + for (auto i = 0; i < rank; i++) { + if (std::find(advanced_ids.begin(), advanced_ids.end(), i) == advanced_ids.end()) { + non_used_dims.push_back(i); + } + } + std::vector permutation_dims; + permutation_dims.insert(permutation_dims.end(), advanced_ids.begin(), advanced_ids.end()); + permutation_dims.insert(permutation_dims.end(), non_used_dims.begin(), non_used_dims.end()); + auto transpose_dims = v0::Constant::create(element::i32, Shape{permutation_dims.size()}, permutation_dims); + auto transposed_input = rg.make(data, transpose_dims); + auto flatten_input = flatten(rg, transposed_input, adv_idx_count); + auto cum_adv_index = masked_indicies[advanced_ids.back()]; + cum_adv_index = rg.make(cum_adv_index, element::i32); + auto multiplier = input_dims->output(advanced_ids.back()); + for (int i = static_cast(adv_idx_count) - 2; i > -1; i--) { + auto input_id = advanced_ids[i]; + auto m_idx = rg.make(masked_indicies[input_id], element::i32); + auto adv_index = rg.make(m_idx, multiplier); + cum_adv_index = rg.make(cum_adv_index, adv_index); + multiplier = rg.make(multiplier, input_dims->output(input_id)); + } + std::shared_ptr gather = rg.make(flatten_input, cum_adv_index, zero); + OutputVector concat_dims; + // check if all advanced indices are consecutive. + std::vector consequence_dims; + auto cum_adv_index_shape_tensor = rg.make(cum_adv_index, element::i32); + for (size_t i = advanced_ids[0]; i <= advanced_ids[advanced_ids.back()]; i++) { + consequence_dims.push_back(i); + } + // unfold regular index axes + if (advanced_ids == consequence_dims) { + OutputVector folded_adv_idx_shape_vector; + auto minus_one = v0::Constant::create(element::i32, Shape{1}, {-1}); + folded_adv_idx_shape_vector.push_back(minus_one); + for (auto i : non_used_dims) { + folded_adv_idx_shape_vector.push_back(input_dims->output(i)); + } + auto folded_adv_idx_shape = rg.make(folded_adv_idx_shape_vector, 0); + gather = rg.make(gather, folded_adv_idx_shape, false); + std::vector adv_idx_permute; + for (size_t i = 1; i < advanced_ids[0] + 1; i++) { + adv_idx_permute.push_back(i); + } + adv_idx_permute.push_back(0); + for (size_t i = advanced_ids[0] + 1; i < (rank - adv_idx_count + 1); i++) { + adv_idx_permute.push_back(i); + } + // Transpose folded advanced indexed axis to its original location. + auto permute_indicies = v0::Constant::create(element::i32, Shape{adv_idx_permute.size()}, adv_idx_permute); + gather = rg.make(gather, permute_indicies); + // unfold advanced index axes + for (size_t i = 0; i < advanced_ids[0]; i++) { + concat_dims.push_back(input_dims->output(i)); + } + concat_dims.push_back(cum_adv_index_shape_tensor); + for (auto i : non_used_dims) { + if (i < advanced_ids[0]) { + continue; + } + concat_dims.push_back(input_dims->output(i)); + } + + } else { + size_t i = 0; + auto one = v0::Constant::create(element::i32, Shape{1}, {1}); + while (i < non_used_dims.size() && non_used_dims[i] < advanced_ids[0]) { + concat_dims.push_back(one); + i++; + } + concat_dims.push_back(cum_adv_index_shape_tensor); + for (; i < non_used_dims.size(); i++) { + concat_dims.push_back(input_dims->output(non_used_dims[i])); + } + } + auto final_shape = rg.make(concat_dims, 0); + gather = rg.make(gather, final_shape, false); + return {gather}; +} +} // namespace + +OutputVector translate_index(const NodeContext& context) { + num_inputs_check(context, 2, 2); + auto x = context.get_input(0); + if (context.input_is_none(1)) { + return {x}; + } + auto indices = context.get_input(1); + auto index_dtype = context.get_input_type(1); + if (index_dtype.is()) { + auto list_elems = get_list_as_outputs(indices); + ov::pass::NodeRegistry rg; + auto rank = x.get_partial_shape().rank(); + // index transformation supports only tensors with static rank + FRONT_END_OP_CONVERSION_CHECK(rank.is_static(), "Dynamic rank for aten::index input is not supported."); + auto res = index_on_list(rg, x, list_elems, rank.get_length()); + context.mark_nodes(rg.get()); + return res; + } + auto index_ov_type = indices.get_element_type(); + if (index_ov_type.is_dynamic()) { + if (simplified_type_interpret(index_dtype).is()) { + index_ov_type = index_dtype.as(); + } + } + if (index_ov_type == element::boolean || index_ov_type == element::u8) { + auto nonzero = context.mark_node(std::make_shared(indices, element::i32)); + auto input_order = context.mark_node(v0::Constant::create(element::i32, Shape{2}, {1, 0})); + auto masked_id = context.mark_node(std::make_shared(nonzero, input_order)); + auto gather = context.mark_node(std::make_shared(x, masked_id)); + return {gather}; + } + if (index_ov_type != element::i32) { + indices = context.mark_node(std::make_shared(indices, element::i32)); + } + auto dim = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); + return {context.mark_node(std::make_shared(x, indices, dim))}; +}; + OutputVector translate_index_fx(const NodeContext& context) { num_inputs_check(context, 2, context.get_input_size()); auto x = context.get_input(0); std::deque> list_elems; for (size_t i = 1; i < context.get_input_size(); i++) { - auto index = context.get_input(static_cast(i)); - if (index.get_element_type() == element::i64) { - auto converted = context.mark_node(std::make_shared(index, element::i32)); - list_elems.push_back(converted); - } else { - list_elems.push_back(index); + Output index; + if (!context.input_is_none(i)) { + index = context.get_input(static_cast(i)); } + list_elems.push_back(index); } - auto concat = - context.mark_node(std::make_shared(OutputVector(list_elems.begin(), list_elems.end()), 0)); - auto gather = std::make_shared(x, concat); - return {gather}; + ov::pass::NodeRegistry rg; + auto rank = x.get_partial_shape().rank(); + if (rank.is_dynamic()) { + rank = context.get_decoder()->get_input_shape(0).rank(); + } + // index transformation supports only tensors with static rank + FRONT_END_OP_CONVERSION_CHECK(rank.is_static(), "Dynamic rank for aten::index input is not supported."); + auto res = index_on_list(rg, x, list_elems, rank.get_length()); + context.mark_nodes(rg.get()); + return res; }; } // namespace op diff --git a/src/frontends/pytorch/src/op/leaky_relu.cpp b/src/frontends/pytorch/src/op/leaky_relu.cpp new file mode 100644 index 00000000000..202f2c82331 --- /dev/null +++ b/src/frontends/pytorch/src/op/leaky_relu.cpp @@ -0,0 +1,32 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/prelu.hpp" +#include "utils.hpp" + +namespace ov { +namespace frontend { +namespace pytorch { +namespace op { + +using namespace ov::op; + +OutputVector translate_leaky_relu_fx(const NodeContext& context) { + num_inputs_check(context, 1, 2); + auto x = context.get_input(0); + float default_negative_slope = 1e-2; + Output negative_slope = ov::op::v0::Constant::create(element::f32, Shape{1}, {default_negative_slope}); + if (context.get_input_size() == 1) { + negative_slope = context.mark_node(std::make_shared(negative_slope, x)); + } else { + negative_slope = context.get_input(1); + } + return {context.mark_node(std::make_shared(x, negative_slope))}; +}; + +} // namespace op +} // namespace pytorch +} // namespace frontend +} // namespace ov diff --git a/src/frontends/pytorch/src/op/mean.cpp b/src/frontends/pytorch/src/op/mean.cpp index 503b7a47f29..ef88c2922ac 100644 --- a/src/frontends/pytorch/src/op/mean.cpp +++ b/src/frontends/pytorch/src/op/mean.cpp @@ -20,7 +20,7 @@ OutputVector translate_mean(const NodeContext& context) { // aten::mean(Tensor self, *, ScalarType? dtype=None) -> Tensor if (num_inputs == 2) { if (!context.input_is_none(1)) { - x = apply_dtype(context, 3, x); + x = apply_dtype(context, 1, x); } axes = get_axes_range(context, 0); } else { @@ -42,6 +42,30 @@ OutputVector translate_mean(const NodeContext& context) { return {mean}; }; +OutputVector translate_mean_fx(const NodeContext& context) { + num_inputs_check(context, 2, 5); + auto x = context.get_input(0); + auto num_inputs = context.get_input_size(); + bool keep_dims = false; + Output axes; + if (num_inputs == 2) { + axes = context.get_input(1); + } else { + axes = context.get_input(1); + if (!context.input_is_none(2)) { + keep_dims = context.const_input(2); + } + if (!context.input_is_none(3)) { + x = apply_dtype(context, 3, x); + } + } + auto mean = context.mark_node(std::make_shared(x, axes, keep_dims)); + if (num_inputs == 5 && !context.input_is_none(4)) { + context.mutate_input(4, mean); + } + return {mean}; +}; + } // namespace op } // namespace pytorch } // namespace frontend diff --git a/src/frontends/pytorch/src/op/pad.cpp b/src/frontends/pytorch/src/op/pad.cpp index ab853fec5f0..390277edcf5 100644 --- a/src/frontends/pytorch/src/op/pad.cpp +++ b/src/frontends/pytorch/src/op/pad.cpp @@ -22,11 +22,12 @@ namespace op { using namespace ov::op; -OutputVector translate_pad(const NodeContext& context) { - num_inputs_check(context, 2, 4); - auto data = context.get_input(0); - auto paddings = context.const_input>(1); - std::string mode = "constant"; +namespace { +OutputVector translate_pad_common(const NodeContext& context, + const Output& data, + const std::vector& paddings, + const Output& pad_value, + const std::string& mode = "constant") { Output shape; Output rank; std::tie(shape, rank) = get_shape_rank(context, data); @@ -45,9 +46,6 @@ OutputVector translate_pad(const NodeContext& context) { auto pads_remaining = context.mark_node(std::make_shared(zero, pads_diff)); auto pads_begins = context.mark_node(std::make_shared(NodeVector{pads_remaining, pads_begin_short}, 0)); auto pads_ends = context.mark_node(std::make_shared(NodeVector{pads_remaining, pads_end_short}, 0)); - if (!context.input_is_none(2)) { - mode = context.const_input(2); - } if (mode == "circular") { int64_t pad_l; int64_t pad_r; @@ -86,24 +84,46 @@ OutputVector translate_pad(const NodeContext& context) { } return {cur}; } + auto pad_value_ = context.mark_node(std::make_shared(pad_value, data)); const std::map pt_to_ov_pad{ {"constant", PadMode::CONSTANT}, {"reflect", PadMode::REFLECT}, {"replicate", PadMode::EDGE}, }; - Output pad_value = context.mark_node(v0::Constant::create(element::f32, Shape{}, {0})); - if (mode == "constant") { - if (!context.input_is_none(3)) { - pad_value = context.get_input(3); - } - pad_value = context.mark_node(std::make_shared(pad_value, data)); - } auto ov_mode = pt_to_ov_pad.find(mode); FRONT_END_OP_CONVERSION_CHECK(ov_mode != pt_to_ov_pad.end(), "aten::pad conversion doesn't support [ ", mode, " ] padding mode"); - return {context.mark_node(std::make_shared(data, pads_begins, pads_ends, pad_value, ov_mode->second))}; + return {context.mark_node(std::make_shared(data, pads_begins, pads_ends, pad_value_, ov_mode->second))}; +} +} // namespace + +OutputVector translate_pad(const NodeContext& context) { + num_inputs_check(context, 2, 4); + auto data = context.get_input(0); + auto paddings = context.const_input>(1); + std::string mode = "constant"; + + if (!context.input_is_none(2)) { + mode = context.const_input(2); + } + + Output pad_value = context.mark_node(v0::Constant::create(element::f32, Shape{}, {0})); + if (mode == "constant") { + if (!context.input_is_none(3)) { + pad_value = context.get_input(3); + } + } + return translate_pad_common(context, data, paddings, pad_value, mode); +} + +OutputVector translate_constant_pad_nd_fx(const NodeContext& context) { + num_inputs_check(context, 3, 3); + auto data = context.get_input(0); + auto paddings = context.const_input>(1); + auto pad_value = context.get_input(2); + return translate_pad_common(context, data, paddings, pad_value); } } // namespace op diff --git a/src/frontends/pytorch/src/op/scalar_tensor.cpp b/src/frontends/pytorch/src/op/scalar_tensor.cpp new file mode 100644 index 00000000000..64428486511 --- /dev/null +++ b/src/frontends/pytorch/src/op/scalar_tensor.cpp @@ -0,0 +1,34 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/convert.hpp" +#include "openvino/op/convert_like.hpp" +#include "pt_framework_node.hpp" +#include "utils.hpp" + +namespace ov { +namespace frontend { +namespace pytorch { +namespace op { + +using namespace ov::op; + +OutputVector translate_scalar_tensor_fx(const NodeContext& context) { + // aten.scalar_tensor.default(-100.0, dtype = torch.float32, layout = torch.strided, device = device(type='cpu')) + num_inputs_check(context, 1, 1); + auto data = context.get_input(0); + if (context.has_attribute("dtype")) { + auto dtype = context.get_attribute("dtype"); + data = context.mark_node(std::make_shared(context.get_input(0), dtype)); + } + // layout and device can be ignored + return {data}; +} + +} // namespace op +} // namespace pytorch +} // namespace frontend +} // namespace ov diff --git a/src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp b/src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp index 5a9e0d1f95c..c1e5fce3abc 100644 --- a/src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp +++ b/src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp @@ -6,6 +6,7 @@ #include "openvino/frontend/pytorch/node_context.hpp" #include "openvino/op/constant.hpp" #include "openvino/op/convert_like.hpp" +#include "openvino/op/unsqueeze.hpp" #include "openvino/op/util/framework_node.hpp" #include "utils.hpp" @@ -45,16 +46,35 @@ OutputVector translate_scaled_dot_product_attention(const NodeContext& context) }; OutputVector translate_scaled_dot_product_attention_fx(const NodeContext& context) { - // aten::scaled_dot_product_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float - // dropout_p=0., bool is_causal=False) TODO: Scale parameter? - num_inputs_check(context, 3, 6); // TODO: Set 7 instead of 6 if `scale` argument supported in FX. - auto output = translate_scaled_dot_product_attention_common(context); - // TODO: scaled_dot_product_flash_attention has 9 outputs but for most cases only - // the first input is used. Rest of the outputs should be returned properly as - // needed. - ov::OutputVector out_vec; - out_vec.push_back(output); - return {context.mark_node(make_list_construct(out_vec))}; + // torch.ops.aten._scaled_dot_product_flash_attention_for_cpu.default(arg1_1, arg2_1, arg3_1, 0.0, True, attn_mask = + // arg0_1, scale = 5.0) + // aten._scaled_dot_product_flash_attention.default(arg0_1, arg1_1, arg2_1, 0.0, True, scale = 5.0) + num_inputs_check(context, 3, 5); + const auto query = context.get_input(0); + const auto key = context.get_input(1); + const auto value = context.get_input(2); + auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); + OutputVector inputs{query, key, value}; + // Index 3 is dropout + auto causal = false; + if (context.has_attribute("is_causal")) { + causal = context.get_attribute("scale"); + } else if (!context.input_is_none(4)) { + causal = context.const_input(4); + } + if (context.has_attribute("attn_mask")) { + const auto attn_mask = context.get_input("attn_mask"); + inputs.push_back(attn_mask); + } else if (context.has_attribute("scale")) { + // if scale exist but attn_mask no, add zero input to fill in gap + inputs.push_back(context.mark_node(std::make_shared(zero, query))); + } + if (context.has_attribute("scale")) { + const auto scale = context.get_input("scale"); + inputs.push_back(context.mark_node(std::make_shared(scale, query))); + } + auto sdpa = context.mark_node(std::make_shared(inputs, causal)); + return {context.mark_node(make_list_construct({sdpa}))}; }; } // namespace op diff --git a/src/frontends/pytorch/src/op/slice_scatter.cpp b/src/frontends/pytorch/src/op/slice_scatter.cpp new file mode 100644 index 00000000000..05bdc081a3e --- /dev/null +++ b/src/frontends/pytorch/src/op/slice_scatter.cpp @@ -0,0 +1,68 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include + +#include "helper_ops/slice_assign.hpp" +#include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/slice.hpp" +#include "openvino/op/unsqueeze.hpp" +#include "utils.hpp" + +namespace ov { +namespace frontend { +namespace pytorch { +namespace op { + +using namespace ov::op; + +OutputVector translate_slice_scatter_fx(const NodeContext& context) { + num_inputs_check(context, 2, 6); + auto input = context.get_input(0); + auto src = context.get_input(1); + auto axis_0 = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); + ov::Output dim; + if (!context.input_is_none(2)) { + dim = context.get_input(2); + if (dim.get_partial_shape().rank().is_dynamic() || dim.get_partial_shape().rank().get_length() == 0) { + dim = context.mark_node(std::make_shared(dim, axis_0)); + } + } else { + dim = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {1})); + } + ov::Output start; + if (!context.input_is_none(3)) { + start = context.get_input(3); + if (start.get_partial_shape().rank().is_dynamic() || start.get_partial_shape().rank().get_length() == 0) { + start = context.mark_node(std::make_shared(start, axis_0)); + } + } else { + start = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0})); + } + ov::Output end; + if (!context.input_is_none(4)) { + end = context.get_input(4); + if (end.get_partial_shape().rank().is_dynamic() || end.get_partial_shape().rank().get_length() == 0) { + end = context.mark_node(std::make_shared(end, axis_0)); + } + } else { + end = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {INT_MAX})); + } + ov::Output step; + if (!context.input_is_none(5)) { + step = context.get_input(5); + if (step.get_partial_shape().rank().is_dynamic() || step.get_partial_shape().rank().get_length() == 0) { + step = context.mark_node(std::make_shared(step, axis_0)); + } + } else { + step = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {1})); + } + return {context.mark_node(std::make_shared(input, src, start, end, step, dim))}; +}; + +} // namespace op +} // namespace pytorch +} // namespace frontend +} // namespace ov \ No newline at end of file diff --git a/src/frontends/pytorch/src/op/split.cpp b/src/frontends/pytorch/src/op/split.cpp index db7d3507580..7b0241192cb 100644 --- a/src/frontends/pytorch/src/op/split.cpp +++ b/src/frontends/pytorch/src/op/split.cpp @@ -4,10 +4,11 @@ #include "openvino/op/split.hpp" -#include +//#include #include "openvino/frontend/pytorch/node_context.hpp" #include "openvino/op/util/framework_node.hpp" +#include "openvino/op/variadic_split.hpp" #include "utils.hpp" namespace ov { @@ -18,7 +19,7 @@ namespace op { using namespace ov::op; OutputVector translate_chunk_fx(const NodeContext& context) { - num_inputs_check(context, 2, 3); + num_inputs_check(context, 3, 3); auto num_chunks = context.const_input(1); auto dim = context.get_input(2); @@ -35,6 +36,17 @@ OutputVector translate_chunk_fx(const NodeContext& context) { return {context.mark_node(make_list_construct(chunk->outputs()))}; } +OutputVector translate_split_with_sizes_fx(const NodeContext& context) { + num_inputs_check(context, 3, 3); + auto data = context.get_input(0); + auto split_lengths = context.get_input(1); + auto dim = context.get_input(2); + + auto split = std::make_shared(data, dim, split_lengths); + + return {context.mark_node(make_list_construct(split->outputs()))}; +} + } // namespace op } // namespace pytorch } // namespace frontend diff --git a/src/frontends/pytorch/src/op/sub.cpp b/src/frontends/pytorch/src/op/sub.cpp index 62534aee538..69148f23329 100644 --- a/src/frontends/pytorch/src/op/sub.cpp +++ b/src/frontends/pytorch/src/op/sub.cpp @@ -21,7 +21,7 @@ OutputVector translate_sub_common(const NodeContext& context, bool inplace) { auto y = context.get_input(1); if (inplace) { if (x.get_element_type().is_dynamic() || x.get_element_type() != y.get_element_type()) - y = context.mark_node(std::make_shared(x, y)); + y = context.mark_node(std::make_shared(y, x)); } else { align_eltwise_input_types(context, x, y); } @@ -45,6 +45,20 @@ OutputVector translate_sub_(const NodeContext& context) { return translate_sub_common(context, true); }; +OutputVector translate_sub_fx(const NodeContext& context) { + num_inputs_check(context, 2, 2); + auto x = context.get_input(0); + auto y = context.get_input(1); + align_eltwise_input_types(context, x, y); + // default alpha is 1 so no need to multiply if alpha is not provided + if (context.has_attribute("alpha")) { + auto alpha = context.get_attribute>("alpha"); + auto casted_alpha = context.mark_node(std::make_shared(alpha, y)); + y = context.mark_node(std::make_shared(casted_alpha, y)); + } + return {context.mark_node(std::make_shared(x, y))}; +}; + } // namespace op } // namespace pytorch } // namespace frontend diff --git a/src/frontends/pytorch/src/op/to.cpp b/src/frontends/pytorch/src/op/to.cpp index 0267e0ed20c..d902caeb848 100644 --- a/src/frontends/pytorch/src/op/to.cpp +++ b/src/frontends/pytorch/src/op/to.cpp @@ -81,6 +81,16 @@ OutputVector translate_to(const NodeContext& context) { return {cast}; } +OutputVector translate_to_fx(const NodeContext& context) { + num_inputs_check(context, 1, 1); + auto data = context.get_input(0); + if (context.has_attribute("dtype")) { + auto dtype = context.get_attribute("dtype"); + data = context.mark_node(std::make_shared(context.get_input(0), dtype)); + } + return {data}; +} + } // namespace op } // namespace pytorch } // namespace frontend diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp index 9bc3083fbfd..c829099c4ac 100644 --- a/src/frontends/pytorch/src/op_table.cpp +++ b/src/frontends/pytorch/src/op_table.cpp @@ -95,6 +95,7 @@ OP_CONVERTER(translate_gru); OP_CONVERTER(translate_hardtanh); OP_CONVERTER(translate_if); OP_CONVERTER(translate_im2col); +OP_CONVERTER(translate_index); OP_CONVERTER(translate_index_add); OP_CONVERTER(translate_index_put_); OP_CONVERTER(translate_index_select); @@ -230,19 +231,36 @@ OP_CONVERTER(translate_quantized_convnd_relu); OP_CONVERTER(translate_quantized_linear); OP_CONVERTER(translate_xor); // Torch FX Translations +OP_CONVERTER(translate_adaptive_max_pool1d_fx); +OP_CONVERTER(translate_adaptive_max_pool2d_fx); +OP_CONVERTER(translate_adaptive_max_pool3d_fx); +OP_CONVERTER(translate_addcmul_fx); +OP_CONVERTER(translate_addmm_fx); OP_CONVERTER(translate_arange_fx); OP_CONVERTER(translate_batch_norm_legit_fx); OP_CONVERTER(translate_batch_norm_legit_no_training_fx); +OP_CONVERTER(translate_batch_norm_legit_no_stats_fx); OP_CONVERTER(translate_cat_fx); +OP_CONVERTER(translate_constant_pad_nd_fx); OP_CONVERTER(translate_chunk_fx); +OP_CONVERTER(translate_div_fx); OP_CONVERTER(translate_expand_fx); +OP_CONVERTER(translate_full_fx); +OP_CONVERTER(translate_gelu_fx); OP_CONVERTER(translate_group_norm_fx); OP_CONVERTER(translate_index_fx); OP_CONVERTER(translate_layer_norm_fx); +OP_CONVERTER(translate_leaky_relu_fx); OP_CONVERTER(translate_max_poolnd_fx); +OP_CONVERTER(translate_mean_fx); +OP_CONVERTER(translate_split_with_sizes_fx); +OP_CONVERTER(translate_scalar_tensor_fx); OP_CONVERTER(translate_scaled_dot_product_attention_fx); OP_CONVERTER(translate_slice_fx); +OP_CONVERTER(translate_slice_scatter_fx); OP_CONVERTER(translate_softmax_fx); +OP_CONVERTER(translate_sub_fx); +OP_CONVERTER(translate_to_fx); OP_CONVERTER(translate_transpose_fx); } // namespace op @@ -631,81 +649,123 @@ const std::map get_supported_ops_ts() { const std::map get_supported_ops_fx() { return { - {"aten._adaptive_avg_pool2d.default", op::translate_1to1_match_2_inputs}, + {"aten.abs.default", op::translate_1to1_match_1_inputs}, + {"aten._adaptive_avg_pool1d.default", op::translate_adaptive_avg_pool1d}, + {"aten._adaptive_avg_pool2d.default", op::translate_adaptive_avg_pool2d}, + {"aten._adaptive_avg_pool3d.default", op::translate_adaptive_avg_pool3d}, + {"aten.adaptive_max_pool1d.default", op::translate_adaptive_max_pool1d_fx}, + {"aten.adaptive_max_pool2d.default", op::translate_adaptive_max_pool2d_fx}, + {"aten.adaptive_max_pool3d.default", op::translate_adaptive_max_pool3d_fx}, + {"aten._local_scalar_dense.default", op::skip_node}, {"aten._softmax.default", op::translate_softmax_fx}, - {"aten._to_copy.default", op::skip_node}, + {"aten._to_copy.default", op::translate_to_fx}, {"aten._unsafe_view.default", op::translate_reshape}, {"aten.add.Tensor", op::translate_add}, {"aten.add_.Tensor", op::translate_add}, - {"aten.addmm.default", op::translate_addmm}, + {"aten.addcmul.default", op::translate_addcmul_fx}, + {"aten.addmm.default", op::translate_addmm_fx}, + {"aten.alias.default", op::skip_node}, {"aten.arange.start", op::translate_arange_fx}, + {"aten.arange.start_step", op::translate_arange_fx}, {"aten.arange.default", op::translate_arange_fx}, {"aten.argmax.default", op::translate_argmax}, + {"aten.as_strided.default", op::translate_as_strided}, {"aten.avg_pool2d.default", op::translate_avg_poolnd}, - {"aten.baddbmm.default", op::translate_addmm}, + {"aten.avg_pool3d.default", op::translate_avg_poolnd}, + {"aten.baddbmm.default", op::translate_addmm_fx}, {"aten.bitwise_and.Tensor", op::translate_bitwise_and}, {"aten.bmm.default", op::translate_1to1_match_2_inputs_align_types}, {"aten.cat.default", op::translate_cat_fx}, + {"aten.ceil.default", op::translate_1to1_match_1_inputs}, + {"aten.clamp.default", op::translate_clamp}, {"aten.clamp_min.default", op::translate_1to1_match_2_inputs}, + {"aten.constant_pad_nd.default", op::translate_constant_pad_nd_fx}, {"aten.clone.default", op::skip_node}, // ignore clone operators that are inserted by PyTorch autograd {"aten.convolution.default", op::translate_convolution}, - {"aten.copy_.default", op::skip_node}, + {"aten._convolution.default", op::translate_convolution}, + {"aten.copy.default", op::skip_node}, + {"aten.copy_.default", op::translate_copy_}, {"aten.cos.default", op::translate_1to1_match_1_inputs}, {"aten.cumsum.default", op::translate_cumsum}, {"aten.detach.default", op::skip_node}, - {"aten.div.Scalar", op::translate_div}, - {"aten.div.Tensor", op::translate_div}, + {"aten.div.Scalar", op::translate_div_fx}, + {"aten.div.Tensor", op::translate_div_fx}, + {"aten.div.Tensor_mode", op::translate_div_fx}, {"aten.embedding.default", op::translate_embedding}, {"aten.empty.memory_format", op::translate_empty}, {"aten.eq.Scalar", op::translate_1to1_match_2_inputs_align_types}, {"aten.eq.Tensor", op::translate_1to1_match_2_inputs_align_types}, {"aten.exp.default", op::translate_1to1_match_1_inputs}, {"aten.expand.default", op::translate_expand_fx}, - {"aten.full.default", op::translate_full}, + {"aten.floor.default", op::translate_1to1_match_1_inputs}, + {"aten.floor_divide.default", op::translate_floor_divide}, + {"aten.full.default", op::translate_full_fx}, + {"aten.full.names", op::translate_full_fx}, + {"aten.full_like.default", op::translate_full_like}, {"aten.gather.default", op::translate_gather}, - {"aten.gelu.default", op::translate_gelu}, + {"aten.gelu.default", op::translate_gelu_fx}, + {"aten.glu.default", op::translate_glu}, {"aten.gt.Scalar", op::translate_1to1_match_2_inputs_align_types}, {"aten.hardsigmoid.default", op::translate_1to1_match_1_inputs}, + {"aten.hardswish.default", op::translate_1to1_match_1_inputs}, {"aten.hardswish_.default", op::inplace_op>}, + {"aten.hardtanh.default", op::translate_hardtanh}, {"aten.hardtanh_.default", op::inplace_op}, {"aten.index.Tensor", op::translate_index_fx}, - {"aten.leaky_relu_.default", op::inplace_op>}, + {"aten.leaky_relu.default", op::translate_leaky_relu_fx}, + {"aten.leaky_relu_.default", op::inplace_op}, {"aten.lift_fresh_copy.default", op::skip_node}, {"aten.linalg_vector_norm.default", op::translate_linalg_vector_norm}, {"aten.log.default", op::translate_log}, {"aten.logsumexp.default", op::translate_logsumexp}, + {"aten.lt.Scalar", op::translate_1to1_match_2_inputs_align_types}, {"aten.lt.Tensor", op::translate_1to1_match_2_inputs_align_types}, {"aten.masked_fill_.Scalar", op::inplace_op}, {"aten.masked_fill.Tensor", op::translate_masked_fill}, {"aten.max_pool2d_with_indices.default", op::translate_max_poolnd_fx}, - {"aten.mean.dim", op::translate_mean}, + {"aten.max_pool3d_with_indices.default", op::translate_max_poolnd_fx}, + {"aten.mean.dim", op::translate_mean_fx}, {"aten.mm.default", op::translate_1to1_match_2_inputs}, {"aten.mul.Tensor", op::translate_1to1_match_2_inputs_align_types}, {"aten.mul.Scalar", op::translate_1to1_match_2_inputs_align_types}, {"aten.native_batch_norm.default", op::translate_batch_norm_legit_fx}, {"aten._native_batch_norm_legit.default", op::translate_batch_norm_legit_fx}, + {"aten._native_batch_norm_legit.no_stats", op::translate_batch_norm_legit_no_stats_fx}, {"aten._native_batch_norm_legit_no_training.default", op::translate_batch_norm_legit_no_training_fx}, + {"aten._native_batch_norm_legit_functional.default", op::translate_batch_norm_legit_fx}, + {"aten.native_dropout.default", op::skip_node}, {"aten.native_group_norm.default", op::translate_group_norm_fx}, {"aten.native_layer_norm.default", op::translate_layer_norm_fx}, + {"aten.ne.Scalar", op::translate_1to1_match_2_inputs_align_types}, {"aten.neg.default", op::translate_neg}, {"aten.new_ones.default", op::translate_new_ones}, {"aten.permute.default", op::translate_1to1_match_2_inputs}, {"aten.pow.Tensor_Scalar", op::translate_pow}, {"aten.relu.default", op::translate_1to1_match_1_inputs}, {"aten.relu_.default", op::inplace_op>}, + {"aten.repeat.default", op::translate_1to1_match_2_inputs}, {"aten.rsub.Scalar", op::translate_rsub}, + {"aten.roll.default", op::translate_roll}, {"aten._scaled_dot_product_flash_attention.default", op::translate_scaled_dot_product_attention_fx}, + {"aten._scaled_dot_product_flash_attention_for_cpu.default", op::translate_scaled_dot_product_attention_fx}, + {"aten.scalar_tensor.default", op::translate_scalar_tensor_fx}, {"aten.select.int", op::translate_select}, {"aten.sigmoid.default", op::translate_1to1_match_1_inputs}, {"aten.silu.default", op::translate_1to1_match_1_inputs}, {"aten.silu_.default", op::inplace_op>}, {"aten.sin.default", op::translate_1to1_match_1_inputs}, {"aten.slice.Tensor", op::translate_slice_fx}, + {"aten.slice_scatter.default", op::translate_slice_scatter_fx}, {"aten.split.Tensor", op::translate_chunk_fx}, - {"aten.sub.default", op::translate_sub}, - {"aten.sub.Tensor", op::translate_sub}, + {"aten.split_with_sizes.default", op::translate_split_with_sizes_fx}, + {"aten.squeeze.dim", op::translate_squeeze}, + {"aten.squeeze.dims", op::translate_squeeze}, + {"aten.sub.default", op::translate_sub_fx}, + {"aten.sub.Tensor", op::translate_sub_fx}, + {"aten.sum.dim_IntList", op::translate_sum}, {"aten.t.default", op::translate_t}, {"aten.tanh.default", op::translate_1to1_match_1_inputs}, + {"aten.unfold.default", op::translate_unfold}, {"aten.transpose.int", op::translate_transpose}, {"aten.unsqueeze.default", op::translate_1to1_match_2_inputs}, {"aten.upsample_nearest2d.default", op::translate_upsample_nearest2d}, diff --git a/src/frontends/pytorch/src/translate_session.cpp b/src/frontends/pytorch/src/translate_session.cpp index 1beed272f43..1791326f41a 100644 --- a/src/frontends/pytorch/src/translate_session.cpp +++ b/src/frontends/pytorch/src/translate_session.cpp @@ -74,10 +74,6 @@ std::shared_ptr TranslateSession::convert_pytorch_model( std::shared_ptr pytorch_model, const TensorMap& external_tensor_map, const std::shared_ptr& input_model) { - // FIXME: don't use global variable to count inlined inputs, no we should use global counter because ID should be - // unique for all new inlined inputs - static size_t inlined_nodes_counter = 100000000; // Suppose there are not graph with more than 10M nodes - std::shared_ptr resulting_model; // define here to make a conversion in a nested scope { auto parameters = std::make_shared(); @@ -128,18 +124,6 @@ std::shared_ptr TranslateSession::convert_pytorch_model( // But this value can be found in the outer scope, for this purpose we create new input for the model to // link with external scope on a higher level. - auto inlined_inputs = node->inlined_inputs(inlined_nodes_counter); - for (size_t i = 0; i < inlined_inputs.size(); ++i) { - size_t fw_tensor_id = inlined_nodes_counter; - ++inlined_nodes_counter; // TODO: Make sure that Decoder side use the same increment for producing ids - if (tensor_map->find(fw_tensor_id) != tensor_map->end()) { - throw std::runtime_error("Duplicated producer for PT value with unique ID: " + - std::to_string(fw_tensor_id) + " produced by inlined_inputs"); - } - - (*tensor_map)[fw_tensor_id] = inlined_inputs[i]; - } - auto raw_inputs = node->inputs(); for (size_t i = 0; i < raw_inputs.size(); ++i) { auto input = raw_inputs.at(i); @@ -286,11 +270,13 @@ OutputVector TranslateSession::convert_node(const NodeContext& context) { if (it != m_translator_map.end()) { return it->second(context); } + OPENVINO_DEBUG << "No translator found for: " << context.get_op_type() << "\n"; } catch (std::exception& e) { exception = e.what(); } catch (...) { exception = "Unknown exception type."; } + OPENVINO_DEBUG << exception << "\n"; try { // Create PtFrameworkNode for everything that wasn't able to be converted normally return make_framework_node(context, exception); @@ -299,6 +285,7 @@ OutputVector TranslateSession::convert_node(const NodeContext& context) { } catch (...) { exception += " Unknown exception happened while creating FrameworkNode with subgraphs"; } + OPENVINO_DEBUG << exception << "\n"; return make_framework_node_ignore_bodies(context, exception); } diff --git a/src/frontends/pytorch/src/utils.cpp b/src/frontends/pytorch/src/utils.cpp index 20f56b9d421..86c6d11e1d6 100644 --- a/src/frontends/pytorch/src/utils.cpp +++ b/src/frontends/pytorch/src/utils.cpp @@ -346,7 +346,7 @@ std::shared_ptr cast_fw_node(std::shared_ptr return fw_node; } -std::shared_ptr make_list_construct(const ov::OutputVector& inputs) { +std::shared_ptr make_list_construct(const ov::OutputVector& inputs) { auto list_construct = std::make_shared<::ov::op::util::FrameworkNode>(inputs, inputs.size()); ov::op::util::FrameworkNodeAttrs attrs; attrs.set_type_name("PTFrameworkNode"); diff --git a/src/frontends/pytorch/src/utils.hpp b/src/frontends/pytorch/src/utils.hpp index 93dc956124a..3cb5798af9f 100644 --- a/src/frontends/pytorch/src/utils.hpp +++ b/src/frontends/pytorch/src/utils.hpp @@ -64,7 +64,7 @@ OutputVector make_framework_node(const NodeContext& context, const std::string& std::shared_ptr cast_fw_node(std::shared_ptr node, const std::string& type); -std::shared_ptr make_list_construct(const ov::OutputVector& inputs); +std::shared_ptr make_list_construct(const ov::OutputVector& inputs); bool is_none_node(const Output& node); @@ -248,8 +248,17 @@ public: virtual bool may_produce_alias(size_t in_index, size_t out_index) const override { FRONT_END_NOT_IMPLEMENTED(may_produce_alias); } - virtual OutputVector inlined_inputs(size_t start_index) const override { - FRONT_END_NOT_IMPLEMENTED(inlined_inputs); + bool is_input_inlined(size_t index) const override { + FRONT_END_NOT_IMPLEMENTED(is_input_inlined); + } + virtual OutputVector inlined_input(size_t index) const override { + FRONT_END_NOT_IMPLEMENTED(inlined_input); + } + virtual ov::Any get_attribute(const std::string& name) const override { + FRONT_END_NOT_IMPLEMENTED(get_attribute); + } + virtual size_t get_named_input(const std::string& name) const override { + FRONT_END_NOT_IMPLEMENTED(get_named_input); } private: diff --git a/tests/layer_tests/pytest.ini b/tests/layer_tests/pytest.ini index db450888077..042dbb6d50a 100644 --- a/tests/layer_tests/pytest.ini +++ b/tests/layer_tests/pytest.ini @@ -5,4 +5,5 @@ markers = precommit_tf_fe precommit_ts_backend precommit_fx_backend + precommit_torch_export timeout diff --git a/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py b/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py index f89af6688ae..3d906447978 100644 --- a/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py +++ b/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py @@ -9,10 +9,12 @@ import os import numpy as np from common.constants import test_device, test_precision from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder +from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder from openvino.frontend import FrontEndManager from openvino.runtime import Core, Type, PartialShape import torch +from packaging import version import openvino.frontend.pytorch.torchdynamo.backend @@ -74,27 +76,59 @@ class PytorchLayerTest: return True return False + def use_torch_export(): + torch_compile_env = os.getenv("PYTORCH_TRACING_MODE") + if torch_compile_env is not None: + return torch_compile_env == "EXPORT" + return False + ov_inputs = flattenize_inputs(inputs) if use_torch_compile_backend(): self.torch_compile_backend_test(model, torch_inputs, custom_eps) else: - trace_model = kwargs.get('trace_model', False) - freeze_model = kwargs.get('freeze_model', True) - with torch.no_grad(): - if kwargs.get('use_convert_model', False): - smodel, converted_model = self.convert_via_mo( - model, torch_inputs, trace_model, dynamic_shapes, ov_inputs, freeze_model) - else: - smodel, converted_model = self.convert_directly_via_frontend( - model, torch_inputs, trace_model, dynamic_shapes, ov_inputs, freeze_model) + if use_torch_export(): + from openvino import convert_model + from torch.export import export + from torch.fx.experimental.proxy_tensor import make_fx - 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( - smodel.inlined_graph, op), f"Operation {op} type doesn't exist in provided graph" + em = export(model, tuple(torch_inputs)) + if version.parse(torch.__version__) >= version.parse("2.2"): + em = em.run_decompositions() + print(em.graph_module.code) + + try: + gm = make_fx(em)(*torch_inputs) + except: + gm = make_fx(em, tracing_mode='symbolic')(*torch_inputs) + + input_shapes = [] + input_types = [] + for input_data in torch_inputs: + input_types.append(input_data.type()) + input_shapes.append(input_data.size()) + + decoder = TorchFXPythonDecoder(gm, gm, input_shapes=input_shapes, input_types=input_types) + converted_model = convert_model(decoder, example_input=torch_inputs) + self._resolve_input_shape_dtype( + converted_model, ov_inputs, dynamic_shapes) + smodel = model + else: + trace_model = kwargs.get('trace_model', False) + freeze_model = kwargs.get('freeze_model', True) + with torch.no_grad(): + if kwargs.get('use_convert_model', False): + smodel, converted_model = self.convert_via_mo( + model, torch_inputs, trace_model, dynamic_shapes, ov_inputs, freeze_model) + else: + smodel, converted_model = self.convert_directly_via_frontend( + model, torch_inputs, trace_model, dynamic_shapes, ov_inputs, freeze_model) + 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( + smodel.inlined_graph, op), f"Operation {op} type doesn't exist in provided graph" # OV infer: core = Core() compiled = core.compile_model(converted_model, ie_device) diff --git a/tests/layer_tests/pytorch_tests/test_adaptive_avg_pool.py b/tests/layer_tests/pytorch_tests/test_adaptive_avg_pool.py index bedd9beb0d4..3b48ef2bfe3 100644 --- a/tests/layer_tests/pytorch_tests/test_adaptive_avg_pool.py +++ b/tests/layer_tests/pytorch_tests/test_adaptive_avg_pool.py @@ -31,6 +31,7 @@ class TestAdaptiveAvgPool3D(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.precommit_ts_backend @pytest.mark.precommit_fx_backend def test_adaptive_avg_pool3d(self, ie_device, precision, ir_version, input_tensor, output_size): @@ -61,6 +62,7 @@ class TestAdaptiveAvgPool2D(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.precommit_ts_backend @pytest.mark.precommit_fx_backend def test_adaptive_avg_pool2d(self, ie_device, precision, ir_version, input_shape, output_size): @@ -91,6 +93,7 @@ class TestAdaptiveAvgPool1D(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.precommit_ts_backend @pytest.mark.precommit_fx_backend def test_adaptive_avg_pool1d(self, ie_device, precision, ir_version, input_shape, output_size): diff --git a/tests/layer_tests/pytorch_tests/test_adaptive_max_pool.py b/tests/layer_tests/pytorch_tests/test_adaptive_max_pool.py index b80a82a4a9d..5ab0da52ff8 100644 --- a/tests/layer_tests/pytorch_tests/test_adaptive_max_pool.py +++ b/tests/layer_tests/pytorch_tests/test_adaptive_max_pool.py @@ -47,6 +47,7 @@ class TestAdaptiveMaxPool3D(PytorchLayerTest): ])) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.precommit_ts_backend @pytest.mark.precommit_fx_backend @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64', @@ -92,6 +93,7 @@ class TestAdaptiveMaxPool2D(PytorchLayerTest): ])) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.precommit_ts_backend @pytest.mark.precommit_fx_backend @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64', @@ -139,6 +141,7 @@ class TestAdaptiveMaxPool1D(PytorchLayerTest): ])) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.precommit_ts_backend @pytest.mark.precommit_fx_backend @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64', diff --git a/tests/layer_tests/pytorch_tests/test_addcmul.py b/tests/layer_tests/pytorch_tests/test_addcmul.py index 18acaf7d201..5ac9aa51b5f 100644 --- a/tests/layer_tests/pytorch_tests/test_addcmul.py +++ b/tests/layer_tests/pytorch_tests/test_addcmul.py @@ -46,6 +46,7 @@ class TestAddCMul(PytorchLayerTest): ]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_addcmul(self, input_type, value, ie_device, precision, ir_version): self.input_type = input_type self._test(*self.create_model(value), ie_device, precision, ir_version) diff --git a/tests/layer_tests/pytorch_tests/test_addmm.py b/tests/layer_tests/pytorch_tests/test_addmm.py index f73894766e8..ac5a78c4fb2 100644 --- a/tests/layer_tests/pytorch_tests/test_addmm.py +++ b/tests/layer_tests/pytorch_tests/test_addmm.py @@ -44,6 +44,7 @@ class TestAddMM(PytorchLayerTest): (1, 1)]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_addmm(self, kwargs_to_prepare_input, alpha, beta, ie_device, precision, ir_version): self._test(*self.create_model(alpha, beta), ie_device, precision, ir_version, kwargs_to_prepare_input=kwargs_to_prepare_input) @@ -93,6 +94,7 @@ class TestBAddBMM(PytorchLayerTest): ]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_baddbmm(self, kwargs_to_prepare_input, alpha, beta, ie_device, precision, ir_version): self._test(*self.create_model(alpha, beta), ie_device, precision, ir_version, kwargs_to_prepare_input=kwargs_to_prepare_input) diff --git a/tests/layer_tests/pytorch_tests/test_and.py b/tests/layer_tests/pytorch_tests/test_and.py index 07e9b1d660d..35b3f51ae37 100644 --- a/tests/layer_tests/pytorch_tests/test_and.py +++ b/tests/layer_tests/pytorch_tests/test_and.py @@ -50,6 +50,7 @@ class TestAnd(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_and_tensor(self, ie_device, precision, ir_version): self.input_data = ( np.array([True, False, False], dtype=np.bool_), @@ -59,18 +60,21 @@ class TestAnd(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_and_bool(self, ie_device, precision, ir_version): self.input_data = (np.array(True, dtype=np.bool_), np.array(True, dtype=np.bool_)) self._test(*self.create_model_bool_input(), ie_device, precision, ir_version) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_and_int(self, ie_device, precision, ir_version): self.input_data = (np.array(3, dtype=np.int32), np.array(4, dtype=np.int32)) self._test(*self.create_model_int_input(), ie_device, precision, ir_version) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_and_tensor(self, ie_device, precision, ir_version): self.input_data = (np.array([3, 5, 8], dtype=np.int32), np.array([7, 11, 2], dtype=np.int32)) self._test( diff --git a/tests/layer_tests/pytorch_tests/test_as_strided.py b/tests/layer_tests/pytorch_tests/test_as_strided.py index 9bfaa66d3a7..fbacb1c81ad 100644 --- a/tests/layer_tests/pytorch_tests/test_as_strided.py +++ b/tests/layer_tests/pytorch_tests/test_as_strided.py @@ -40,6 +40,7 @@ class TestAsStrided(PytorchLayerTest): @pytest.mark.parametrize("offset", [None, 1, 3, 7]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_as_strided(self, size, stride, offset, ie_device, precision, ir_version): self._test(*self.create_model(size, stride, offset), ie_device, precision, ir_version, trace_model=True) @@ -90,6 +91,7 @@ class TestAsStridedListConstruct(PytorchLayerTest): @pytest.mark.parametrize("mode", ["no_const", "stride_const", "size_const"]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_as_strided_list_construct(self, size, stride, offset, mode, ie_device, precision, ir_version): inp_kwargs = {"size_shape_tensor": size, "stride_shape_tensor": stride} self._test( @@ -121,5 +123,6 @@ class TestAsStridedLongformer(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_as_strided_lf(self, ie_device, precision, ir_version): self._test(*self.create_model(), ie_device, precision, ir_version, trace_model=True, freeze_model=False) diff --git a/tests/layer_tests/pytorch_tests/test_batch_norm.py b/tests/layer_tests/pytorch_tests/test_batch_norm.py index 94208c7359a..2275f53960a 100644 --- a/tests/layer_tests/pytorch_tests/test_batch_norm.py +++ b/tests/layer_tests/pytorch_tests/test_batch_norm.py @@ -59,6 +59,7 @@ class TestBatchNorm(PytorchLayerTest): @pytest.mark.precommit @pytest.mark.precommit_ts_backend @pytest.mark.precommit_fx_backend + @pytest.mark.precommit_torch_export 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) \ No newline at end of file diff --git a/tests/layer_tests/pytorch_tests/test_broadcast_tensors.py b/tests/layer_tests/pytorch_tests/test_broadcast_tensors.py index b405f0114b5..2042dde0c46 100644 --- a/tests/layer_tests/pytorch_tests/test_broadcast_tensors.py +++ b/tests/layer_tests/pytorch_tests/test_broadcast_tensors.py @@ -31,6 +31,7 @@ class TestBroadcastTensors(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.parametrize("x_shape", [[1, ], [2, 1], [2, 2, 1]]) @pytest.mark.parametrize("y_shape", [[2, ], [1, 2], [1, 2, 1]]) @pytest.mark.parametrize("z_shape", [[1, 2], [2, 2], [1, 2, 1, 1]]) diff --git a/tests/layer_tests/pytorch_tests/test_clone.py b/tests/layer_tests/pytorch_tests/test_clone.py index f56aac44e9d..18e07d0efd0 100644 --- a/tests/layer_tests/pytorch_tests/test_clone.py +++ b/tests/layer_tests/pytorch_tests/test_clone.py @@ -25,5 +25,6 @@ class TestClone(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_clone(self, ie_device, precision, ir_version): self._test(*self.create_model(), ie_device, precision, ir_version) diff --git a/tests/layer_tests/pytorch_tests/test_conv_transposend.py b/tests/layer_tests/pytorch_tests/test_conv_transposend.py index db0350d649e..7b322918859 100644 --- a/tests/layer_tests/pytorch_tests/test_conv_transposend.py +++ b/tests/layer_tests/pytorch_tests/test_conv_transposend.py @@ -62,6 +62,7 @@ class TestConvTranspose2D(PytorchLayerTest): @pytest.mark.parametrize("bias", [True, False]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_conv_transpose2d(self, params, bias, ie_device, precision, ir_version): self._test(*self.create_model(**params, bias=bias), ie_device, precision, ir_version, dynamic_shapes=params['groups'] == 1) @@ -121,6 +122,7 @@ class TestConvTranspose1D(PytorchLayerTest): @pytest.mark.parametrize("bias", [True, False]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_conv_transpose1d(self, params, bias, ie_device, precision, ir_version): self._test(*self.create_model(**params, bias=bias), ie_device, precision, ir_version, dynamic_shapes=params['groups'] == 1) @@ -191,6 +193,7 @@ class TestConvTranspose3D(PytorchLayerTest): @pytest.mark.parametrize("bias", [True, False]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_conv_transpose3d(self, params, bias, ie_device, precision, ir_version): self._test(*self.create_model(**params, bias=bias), ie_device, precision, ir_version, dynamic_shapes=params['groups'] == 1) diff --git a/tests/layer_tests/pytorch_tests/test_convolution.py b/tests/layer_tests/pytorch_tests/test_convolution.py index 5ac994f9d79..19e4c9b7cee 100644 --- a/tests/layer_tests/pytorch_tests/test_convolution.py +++ b/tests/layer_tests/pytorch_tests/test_convolution.py @@ -209,6 +209,7 @@ class TestConvolution(PytorchLayerTest): @pytest.mark.precommit @pytest.mark.precommit_ts_backend @pytest.mark.precommit_fx_backend + @pytest.mark.precommit_torch_export 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, @@ -221,6 +222,7 @@ class TestConvolution(PytorchLayerTest): @pytest.mark.precommit @pytest.mark.precommit_ts_backend @pytest.mark.precommit_fx_backend + @pytest.mark.precommit_torch_export 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) diff --git a/tests/layer_tests/pytorch_tests/test_device.py b/tests/layer_tests/pytorch_tests/test_device.py index 2efdf5c85c5..72982f7a677 100644 --- a/tests/layer_tests/pytorch_tests/test_device.py +++ b/tests/layer_tests/pytorch_tests/test_device.py @@ -50,6 +50,7 @@ class TestDevice(PytorchLayerTest): @pytest.mark.parametrize("device_string", ["cpu", "cuda"]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_device(self, device_string, ie_device, precision, ir_version): self._test( *self.create_model_device(device_string), @@ -63,6 +64,7 @@ class TestDevice(PytorchLayerTest): @pytest.mark.parametrize("device_string", ["cpu", "cuda"]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_device_type(self, device_string, ie_device, precision, ir_version): self._test( *self.create_model_type(device_string), diff --git a/tests/layer_tests/pytorch_tests/test_div.py b/tests/layer_tests/pytorch_tests/test_div.py index 22243e1a591..a2809b1fa68 100644 --- a/tests/layer_tests/pytorch_tests/test_div.py +++ b/tests/layer_tests/pytorch_tests/test_div.py @@ -43,6 +43,7 @@ class TestDiv(PytorchLayerTest): ])) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_div_pt_spec(self, input_array, other_array, rounding_mode, ie_device, precision, ir_version): self.input_array = input_array self.input_type = np.float32 @@ -118,6 +119,7 @@ class TestDivTypes(PytorchLayerTest): ])) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.xfail(condition=platform.system() in ('Darwin', 'Linux') and platform.machine() in ('arm', 'armv7l', 'aarch64', 'arm64', 'ARM64'), diff --git a/tests/layer_tests/pytorch_tests/test_embedding.py b/tests/layer_tests/pytorch_tests/test_embedding.py index a448e21d3ff..901f72d4d76 100644 --- a/tests/layer_tests/pytorch_tests/test_embedding.py +++ b/tests/layer_tests/pytorch_tests/test_embedding.py @@ -26,6 +26,7 @@ class TestEmbedding(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.parametrize("indicies_size", [1, 2, 3, 4]) @pytest.mark.parametrize("indicies_dtype", ["int", "int32"]) def test_embedding(self, ie_device, precision, ir_version, indicies_size, indicies_dtype): diff --git a/tests/layer_tests/pytorch_tests/test_eq.py b/tests/layer_tests/pytorch_tests/test_eq.py index 6fa5f4ecfba..1a73084359b 100644 --- a/tests/layer_tests/pytorch_tests/test_eq.py +++ b/tests/layer_tests/pytorch_tests/test_eq.py @@ -40,6 +40,7 @@ class TestEq(PytorchLayerTest): ]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_eq_pt_spec(self, input_array, other_array, types, ie_device, precision, ir_version): self.input_array = input_array self.input_type = types[0] diff --git a/tests/layer_tests/pytorch_tests/test_expand.py b/tests/layer_tests/pytorch_tests/test_expand.py index 5848aca8fa8..63453fce9d1 100644 --- a/tests/layer_tests/pytorch_tests/test_expand.py +++ b/tests/layer_tests/pytorch_tests/test_expand.py @@ -35,6 +35,7 @@ class TestExpand(PytorchLayerTest): @pytest.mark.parametrize("op_type", ["expand", "broadcast_to"]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_expand(self, dims, op_type, ie_device, precision, ir_version): self._test(*self.create_model(dims, op_type), ie_device, precision, ir_version) @@ -68,6 +69,7 @@ class TestExpandList(PytorchLayerTest): @pytest.mark.parametrize("op_type", ["expand", "broadcast_to"]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_expand(self, dims, op_type, ie_device, precision, ir_version): self._test(*self.create_model(op_type), ie_device, precision, ir_version, kwargs_to_prepare_input={"broadcast_shape": dims}) @@ -104,6 +106,7 @@ class TestExpandAs(PytorchLayerTest): ]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_expand(self, ie_device, precision, ir_version, kwargs_to_prepare_input): self._test(*self.create_model(), ie_device, precision, ir_version, kwargs_to_prepare_input=kwargs_to_prepare_input) diff --git a/tests/layer_tests/pytorch_tests/test_flatten.py b/tests/layer_tests/pytorch_tests/test_flatten.py index 285c86d812c..c6969623325 100644 --- a/tests/layer_tests/pytorch_tests/test_flatten.py +++ b/tests/layer_tests/pytorch_tests/test_flatten.py @@ -37,6 +37,7 @@ class TestFlatten(PytorchLayerTest): [2, 3]]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_flatten(self, dim0, dim1, ie_device, precision, ir_version): self._test(*self.create_model(dim0, dim1), ie_device, precision, ir_version) diff --git a/tests/layer_tests/pytorch_tests/test_floor_divide.py b/tests/layer_tests/pytorch_tests/test_floor_divide.py index 40d1eb5afa7..bd8827b2588 100644 --- a/tests/layer_tests/pytorch_tests/test_floor_divide.py +++ b/tests/layer_tests/pytorch_tests/test_floor_divide.py @@ -51,6 +51,7 @@ class TestFloorDivide(PytorchLayerTest): ])) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64', reason='Ticket - 122715') def test_floor_divide(self, input_tensor, other_tensor, ie_device, precision, ir_version): @@ -79,6 +80,7 @@ class TestFloorDivide(PytorchLayerTest): ]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_floor_divide_int(self, input_data, other_data, ie_device, precision, ir_version): input_tensor = input_data["tensor"] if type(input_tensor) is list: diff --git a/tests/layer_tests/pytorch_tests/test_gelu.py b/tests/layer_tests/pytorch_tests/test_gelu.py index 87177c4cde6..6a6d5e61a57 100644 --- a/tests/layer_tests/pytorch_tests/test_gelu.py +++ b/tests/layer_tests/pytorch_tests/test_gelu.py @@ -29,6 +29,7 @@ class TestGelu(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.parametrize("approximate", ["none", "tanh"]) def test_glu(self, approximate, ie_device, precision, ir_version): self._test(*self.create_model(approximate), ie_device, precision, ir_version) \ No newline at end of file diff --git a/tests/layer_tests/pytorch_tests/test_glu.py b/tests/layer_tests/pytorch_tests/test_glu.py index 718604abc7b..aa77fb7d776 100644 --- a/tests/layer_tests/pytorch_tests/test_glu.py +++ b/tests/layer_tests/pytorch_tests/test_glu.py @@ -29,6 +29,7 @@ class TestGlu(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.parametrize("dim", [0, 1, 2, 3, -1, -2]) def test_glu(self, dim, ie_device, precision, ir_version): self._test(*self.create_model(dim), ie_device, precision, ir_version) \ No newline at end of file diff --git a/tests/layer_tests/pytorch_tests/test_group_norm.py b/tests/layer_tests/pytorch_tests/test_group_norm.py index effb8ea9048..347d1e013a3 100644 --- a/tests/layer_tests/pytorch_tests/test_group_norm.py +++ b/tests/layer_tests/pytorch_tests/test_group_norm.py @@ -53,6 +53,7 @@ class TestGroupNorm(PytorchLayerTest): ]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_group_norm(self, params, ie_device, precision, ir_version, kwargs_to_prepare_input): self._test(*self.create_model(**params), ie_device, precision, ir_version, kwargs_to_prepare_input=kwargs_to_prepare_input) diff --git a/tests/layer_tests/pytorch_tests/test_index.py b/tests/layer_tests/pytorch_tests/test_index.py index 4b4e53690f9..d7f3a18469a 100644 --- a/tests/layer_tests/pytorch_tests/test_index.py +++ b/tests/layer_tests/pytorch_tests/test_index.py @@ -3,6 +3,7 @@ import pytest import numpy as np +import torch from pytorch_layer_test_class import PytorchLayerTest @@ -115,6 +116,7 @@ class TestIndexRange(PytorchLayerTest): def test_index_range(self, input_shape, idx, ie_device, precision, ir_version): self._test(*self.create_model(), ie_device, precision, ir_version, kwargs_to_prepare_input={ "input_shape": input_shape, "idx": idx}, trace_model=True, dynamic_shapes=False) + @pytest.mark.nightly @pytest.mark.precommit @pytest.mark.parametrize(("input_shape", "idx"), ( @@ -126,6 +128,7 @@ class TestIndexRange(PytorchLayerTest): self._test(*self.create_model2(), ie_device, precision, ir_version, kwargs_to_prepare_input={ "input_shape": input_shape, "idx": idx}, trace_model=True, dynamic_shapes=False) + class TestIndexMask(PytorchLayerTest): def _prepare_input(self, input_shape): import numpy as np @@ -151,3 +154,27 @@ class TestIndexMask(PytorchLayerTest): def test_index_mask(self, input_shape, ie_device, precision, ir_version): self._test(*self.create_model(), ie_device, precision, ir_version, kwargs_to_prepare_input={ "input_shape": input_shape}, trace_model=True, use_convert_model=True) + + +class TestIndexNone(PytorchLayerTest): + def _prepare_input(self, input_shape): + import numpy as np + return (np.random.randn(*input_shape).astype(np.float32),) + + class aten_index_list(torch.nn.Module): + def __init__(self, idxs): + super(TestIndexNone.aten_index_list, self).__init__() + self.idxs = idxs + + def forward(self, x): + return x[self.idxs] + + @pytest.mark.nightly + @pytest.mark.parametrize(("input_shape,idxs"), [ + ((2, 3, 4, 5), (torch.unsqueeze(torch.randint(0, 2, [14], dtype=torch.int32), 1),)), + ((2, 3, 4, 5), (torch.unsqueeze(torch.randint(0, 2, [14], dtype=torch.int32), 1), torch.randint(0, 3, [14], dtype=torch.int32))), + ((2, 3, 4, 5), (None, None, torch.unsqueeze(torch.randint(0, 2, [14], dtype=torch.int32), 1), torch.randint(0, 3, [14], dtype=torch.int32))), + ]) + def test_index(self, input_shape, idxs, ie_device, precision, ir_version): + self._test(self.aten_index_list(idxs), None, "aten::index", ie_device, precision, + ir_version,kwargs_to_prepare_input={"input_shape": input_shape}, use_convert_model=True, trace_model=True) diff --git a/tests/layer_tests/pytorch_tests/test_mm.py b/tests/layer_tests/pytorch_tests/test_mm.py index b0f920e47f5..83de1188d11 100644 --- a/tests/layer_tests/pytorch_tests/test_mm.py +++ b/tests/layer_tests/pytorch_tests/test_mm.py @@ -41,6 +41,7 @@ class TestMatMul(PytorchLayerTest): ]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_mm(self, kwargs_to_prepare_input, ie_device, precision, ir_version): self._test(*self.create_model('aten::mm'), ie_device, precision, ir_version, kwargs_to_prepare_input=kwargs_to_prepare_input) @@ -55,6 +56,7 @@ class TestMatMul(PytorchLayerTest): ]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_bmm(self, kwargs_to_prepare_input, ie_device, precision, ir_version): self._test(*self.create_model('aten::bmm'), ie_device, precision, ir_version, kwargs_to_prepare_input=kwargs_to_prepare_input) @@ -81,6 +83,7 @@ class TestMatMul(PytorchLayerTest): ]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_matmul(self, kwargs_to_prepare_input, ie_device, precision, ir_version): self._test(*self.create_model('aten::matmul'), ie_device, precision, ir_version, kwargs_to_prepare_input=kwargs_to_prepare_input) diff --git a/tests/layer_tests/pytorch_tests/test_mul.py b/tests/layer_tests/pytorch_tests/test_mul.py index c7aea00ee8a..1eca86c6e17 100644 --- a/tests/layer_tests/pytorch_tests/test_mul.py +++ b/tests/layer_tests/pytorch_tests/test_mul.py @@ -32,6 +32,7 @@ class TestMul(PytorchLayerTest): ]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_mul_pt_spec(self, input_array, other_array, ie_device, precision, ir_version): self.input_array = input_array self.input_type = np.float32 @@ -100,6 +101,7 @@ class TestMulTypes(PytorchLayerTest): ]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_mul_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 diff --git a/tests/layer_tests/pytorch_tests/test_permute.py b/tests/layer_tests/pytorch_tests/test_permute.py index b0efdb72a5d..f5a234fc419 100644 --- a/tests/layer_tests/pytorch_tests/test_permute.py +++ b/tests/layer_tests/pytorch_tests/test_permute.py @@ -29,6 +29,7 @@ class TestPermute(PytorchLayerTest): @pytest.mark.parametrize("order", [[0, 2, 3, 1], [0, 3, 1, 2]]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_permute(self, order, ie_device, precision, ir_version): self._test(*self.create_model(order), ie_device, precision, ir_version) @@ -53,5 +54,6 @@ class TestPermuteList(PytorchLayerTest): @pytest.mark.parametrize("order", [[1, 3, 4, 2], [1, 4, 2, 3]]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_permute(self, order, ie_device, precision, ir_version): self._test(*self.create_model(), ie_device, precision, ir_version, kwargs_to_prepare_input={"permute_shape": order}) diff --git a/tests/layer_tests/pytorch_tests/test_reshape.py b/tests/layer_tests/pytorch_tests/test_reshape.py index 5b89cab94d0..f093c8e3bcb 100644 --- a/tests/layer_tests/pytorch_tests/test_reshape.py +++ b/tests/layer_tests/pytorch_tests/test_reshape.py @@ -38,5 +38,6 @@ class TestReshape(PytorchLayerTest): ]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_reshape(self, shape, ie_device, precision, ir_version): self._test(*self.create_model(shape), ie_device, precision, ir_version) diff --git a/tests/layer_tests/pytorch_tests/test_reshape_as.py b/tests/layer_tests/pytorch_tests/test_reshape_as.py index d354a530884..1d5c00cc05f 100644 --- a/tests/layer_tests/pytorch_tests/test_reshape_as.py +++ b/tests/layer_tests/pytorch_tests/test_reshape_as.py @@ -32,6 +32,7 @@ class TestReshapeAs(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.parametrize("op", ["reshape_as", "view_as"]) @pytest.mark.parametrize('input_tensor_shapes',( ((3, 6), (2, 9)), ((2, 2, 3), (6, 2)), ((6, 2), (2, 2, 3)))) def test_reshape_as(self, op, input_tensor_shapes, ie_device, precision, ir_version): diff --git a/tests/layer_tests/pytorch_tests/test_resolve_conj_neg.py b/tests/layer_tests/pytorch_tests/test_resolve_conj_neg.py index 87097b88af3..b26173e94a2 100644 --- a/tests/layer_tests/pytorch_tests/test_resolve_conj_neg.py +++ b/tests/layer_tests/pytorch_tests/test_resolve_conj_neg.py @@ -40,6 +40,7 @@ class TestResolveConjNeg(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.parametrize("op_type", ["resolve_neg", "resolve_conj"]) @pytest.mark.parametrize("dtype", ["float32", "int32"]) def test_reslove(self, op_type, dtype, ie_device, precision, ir_version): @@ -47,6 +48,7 @@ class TestResolveConjNeg(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export @pytest.mark.parametrize("op_type", ["resolve_neg", "resolve_conj"]) @pytest.mark.xfail(reason="complex dtype is not supported yet") def test_resolve_complex(self, op_type, ie_device, precision, ir_version): diff --git a/tests/layer_tests/pytorch_tests/test_scaled_dot_product_attention.py b/tests/layer_tests/pytorch_tests/test_scaled_dot_product_attention.py index d442b9e86a4..94e7b43be85 100644 --- a/tests/layer_tests/pytorch_tests/test_scaled_dot_product_attention.py +++ b/tests/layer_tests/pytorch_tests/test_scaled_dot_product_attention.py @@ -12,13 +12,13 @@ class TestScaledDotProductAttention(PytorchLayerTest): def _prepare_input(self, dtype): return (np.random.randn(1, 2, 8, 4).astype(dtype), np.random.randn(1, 2, 8, 4).astype(dtype), np.random.randn(1, 2, 8, 4).astype(dtype)) - def create_model(self, mask, is_causal, scale, dtype): + def create_model(self, mask, is_causal, dtype): import torch.nn.functional as F import torch class aten_scaled_dot_product_atten(torch.nn.Module): - def __init__(self, mask=False, is_causal=False, scale=False, dtype=np.float32) -> None: + def __init__(self, mask=False, is_causal=False, dtype=np.float32) -> None: super().__init__() self.mask = None if not mask else torch.from_numpy( @@ -28,22 +28,20 @@ class TestScaledDotProductAttention(PytorchLayerTest): self.mask.to(torch.bool) self.is_causal = False - self.scale = None if not scale else torch.tensor( - 5, dtype=torch.float) - def forward(self, query, key, value): - return F.scaled_dot_product_attention(query, key, value, attn_mask=self.mask, is_causal=self.is_causal, scale=self.scale) + # torch export struggles with dynamic scale + a = F.scaled_dot_product_attention(query, key, value, attn_mask=self.mask, is_causal=self.is_causal) + b = F.scaled_dot_product_attention(query, key, value, attn_mask=self.mask, is_causal=self.is_causal, scale=torch.tensor(5, dtype=torch.float)) + return a, b - ref_net = None - - return aten_scaled_dot_product_atten(mask, is_causal, scale, dtype), ref_net, 'aten::scaled_dot_product_attention' + return aten_scaled_dot_product_atten(mask, is_causal, dtype), None, 'aten::scaled_dot_product_attention' @pytest.mark.nightly @pytest.mark.precommit @pytest.mark.precommit_fx_backend + @pytest.mark.precommit_torch_export @pytest.mark.parametrize(['mask', 'is_causal'], [(False, False), (False, True), (True, True), (True, False)]) - @pytest.mark.parametrize('scale', [False, True]) @pytest.mark.parametrize("dtype", (np.float32, np.float64)) - def test_scaled_dot_product_atten(self, ie_device, precision, ir_version, mask, is_causal, scale, dtype): - self._test(*self.create_model(mask, is_causal, scale, dtype), + def test_scaled_dot_product_atten(self, ie_device, precision, ir_version, mask, is_causal, dtype): + self._test(*self.create_model(mask, is_causal, dtype), ie_device, precision, ir_version, kwargs_to_prepare_input={"dtype": dtype}) diff --git a/tests/layer_tests/pytorch_tests/test_select.py b/tests/layer_tests/pytorch_tests/test_select.py index fd1c325d4a6..e11f4a6f102 100644 --- a/tests/layer_tests/pytorch_tests/test_select.py +++ b/tests/layer_tests/pytorch_tests/test_select.py @@ -32,6 +32,7 @@ class TestSelect(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_select(self, ie_device, precision, ir_version, input_dim, input_index): self._test(*self.create_model(input_dim, input_index), ie_device, precision, ir_version) diff --git a/tests/layer_tests/pytorch_tests/test_silu.py b/tests/layer_tests/pytorch_tests/test_silu.py index 4477184dba7..3c206b073de 100644 --- a/tests/layer_tests/pytorch_tests/test_silu.py +++ b/tests/layer_tests/pytorch_tests/test_silu.py @@ -28,5 +28,6 @@ class TestSilu(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_silu(self, ie_device, precision, ir_version): self._test(*self.create_model(), ie_device, precision, ir_version) diff --git a/tests/layer_tests/pytorch_tests/test_softmax.py b/tests/layer_tests/pytorch_tests/test_softmax.py index eac3f9ba3c6..820e99fad17 100644 --- a/tests/layer_tests/pytorch_tests/test_softmax.py +++ b/tests/layer_tests/pytorch_tests/test_softmax.py @@ -49,6 +49,7 @@ class TestSoftmax(PytorchLayerTest): @pytest.mark.parametrize("dim", [-1, 3]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_softmax(self, dim, ie_device, precision, ir_version): self._test(*self.create_model(dim), ie_device, precision, ir_version) @@ -57,6 +58,7 @@ class TestSoftmax(PytorchLayerTest): @pytest.mark.parametrize("use_prim_dtype", [True, False]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_softmax(self, dim, dtype, use_prim_dtype, ie_device, precision, ir_version): input_kwargs = {} if use_prim_dtype: diff --git a/tests/layer_tests/pytorch_tests/test_sub.py b/tests/layer_tests/pytorch_tests/test_sub.py index 5ba3f1a5506..e24a878239e 100644 --- a/tests/layer_tests/pytorch_tests/test_sub.py +++ b/tests/layer_tests/pytorch_tests/test_sub.py @@ -110,6 +110,7 @@ class TestSubTypes(PytorchLayerTest): ]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_sub_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 diff --git a/tests/layer_tests/pytorch_tests/test_tensor_split.py b/tests/layer_tests/pytorch_tests/test_tensor_split.py index 168cbd6f60f..eeeea7542fb 100644 --- a/tests/layer_tests/pytorch_tests/test_tensor_split.py +++ b/tests/layer_tests/pytorch_tests/test_tensor_split.py @@ -61,6 +61,7 @@ class TestTensorSplit(PytorchLayerTest): @pytest.mark.parametrize("axis", [0, 1, -1]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_tensor_split(self, input_shape, splits, axis, ie_device, precision, ir_version): self.input_shape = input_shape self._test(*self.create_model(splits, axis), ie_device, precision, ir_version) diff --git a/tests/layer_tests/pytorch_tests/test_type_as.py b/tests/layer_tests/pytorch_tests/test_type_as.py index 35fe0732602..3f520305794 100644 --- a/tests/layer_tests/pytorch_tests/test_type_as.py +++ b/tests/layer_tests/pytorch_tests/test_type_as.py @@ -28,6 +28,7 @@ class TestTypeAs(PytorchLayerTest): @pytest.mark.parametrize("cast_dtype", [np.float64, np.float32, np.int64, np.int32, np.int16, np.int8, np.uint8]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_type_as(self, input_dtype, cast_dtype, ie_device, precision, ir_version): self._test(*self.create_model(), ie_device, precision, ir_version, kwargs_to_prepare_input={"input_dtype": input_dtype, "cast_dtype": cast_dtype}) diff --git a/tests/layer_tests/pytorch_tests/test_unflatten.py b/tests/layer_tests/pytorch_tests/test_unflatten.py index 3f8e9de3a2b..0c9537912bf 100644 --- a/tests/layer_tests/pytorch_tests/test_unflatten.py +++ b/tests/layer_tests/pytorch_tests/test_unflatten.py @@ -31,6 +31,7 @@ class TestUnflatten(PytorchLayerTest): @pytest.mark.parametrize("dtype", ["float32", "int32"]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_unflatten(self, dim, shape, dtype, ie_device, precision, ir_version): self._test(*self.create_model(dim, shape), ie_device, precision, ir_version, kwargs_to_prepare_input={"dtype": dtype}) @@ -65,5 +66,6 @@ class TestUnflattenListSizes(PytorchLayerTest): @pytest.mark.parametrize("dtype", ["float32", "int32"]) @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_unflatten(self, dim, dtype, ie_device, precision, ir_version): self._test(*self.create_model(dim), ie_device, precision, ir_version, kwargs_to_prepare_input={"dtype": dtype}) \ No newline at end of file diff --git a/tests/layer_tests/pytorch_tests/test_unfold.py b/tests/layer_tests/pytorch_tests/test_unfold.py index b4138fa772b..37a0e467544 100644 --- a/tests/layer_tests/pytorch_tests/test_unfold.py +++ b/tests/layer_tests/pytorch_tests/test_unfold.py @@ -38,6 +38,7 @@ class TestUnfold(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit + @pytest.mark.precommit_torch_export def test_unfold(self, ie_device, precision, ir_version, dimension, size, step, input_shape): self.input_tensor = np.random.randn(*input_shape).astype(np.float32) self._test(*self.create_model(dimension, size, step), diff --git a/tests/model_hub_tests/torch_tests/requirements.txt b/tests/model_hub_tests/torch_tests/requirements.txt index af6f4cf2212..b734007dafa 100644 --- a/tests/model_hub_tests/torch_tests/requirements.txt +++ b/tests/model_hub_tests/torch_tests/requirements.txt @@ -7,6 +7,7 @@ datasets facexlib numpy optimum +packaging pandas protobuf pyctcdecode diff --git a/tests/model_hub_tests/torch_tests/test_aliked.py b/tests/model_hub_tests/torch_tests/test_aliked.py index ad29187de06..a2e93c7ce3c 100644 --- a/tests/model_hub_tests/torch_tests/test_aliked.py +++ b/tests/model_hub_tests/torch_tests/test_aliked.py @@ -84,7 +84,7 @@ class TestAlikedConvertModel(TestTorchConvertModel): subprocess.check_call(["sh", "build.sh"], cwd=os.path.join( self.repo_dir.name, "custom_ops")) - def load_model(self, model_name, model_link): + def load_model_impl(self, model_name, model_link): sys.path.append(self.repo_dir.name) from nets.aliked import ALIKED diff --git a/tests/model_hub_tests/torch_tests/test_detectron2.py b/tests/model_hub_tests/torch_tests/test_detectron2.py index 5c2ea9a1d9e..0c6c9eddb0f 100644 --- a/tests/model_hub_tests/torch_tests/test_detectron2.py +++ b/tests/model_hub_tests/torch_tests/test_detectron2.py @@ -23,7 +23,7 @@ class TestDetectron2ConvertModel(TestTorchConvertModel): subprocess.run([sys.executable, "-m", "pip", "install", "git+https://github.com/facebookresearch/detectron2.git@017abbfa5f2c2a2afa045200c2af9ccf2fc6227f"]) - def load_model(self, model_name, model_link): + def load_model_impl(self, model_name, model_link): from detectron2 import model_zoo, export from detectron2.modeling import build_model, PanopticFPN from detectron2.checkpoint import DetectionCheckpointer diff --git a/tests/model_hub_tests/torch_tests/test_edsr.py b/tests/model_hub_tests/torch_tests/test_edsr.py index 9d4a9d183e8..960af9fd558 100644 --- a/tests/model_hub_tests/torch_tests/test_edsr.py +++ b/tests/model_hub_tests/torch_tests/test_edsr.py @@ -38,7 +38,7 @@ torch.manual_seed(0) class TestEdsrConvertModel(TestTorchConvertModel): - def load_model(self, model_name, model_link): + def load_model_impl(self, model_name, model_link): # image link from https://github.com/eugenesiow/super-image url = 'https://paperswithcode.com/media/datasets/Set5-0000002728-07a9793f_zA3bDjj.jpg' image = Image.open(requests.get(url, stream=True).raw) diff --git a/tests/model_hub_tests/torch_tests/test_gfpgan.py b/tests/model_hub_tests/torch_tests/test_gfpgan.py index 90784926183..ca22789023b 100644 --- a/tests/model_hub_tests/torch_tests/test_gfpgan.py +++ b/tests/model_hub_tests/torch_tests/test_gfpgan.py @@ -28,7 +28,7 @@ class TestGFPGANConvertModel(TestTorchConvertModel): subprocess.check_call( ["wget", "-nv", checkpoint_url], cwd=self.repo_dir.name) - def load_model(self, model_name, model_link): + def load_model_impl(self, model_name, model_link): sys.path.append(self.repo_dir.name) from gfpgan import GFPGANer diff --git a/tests/model_hub_tests/torch_tests/test_hf_transformers.py b/tests/model_hub_tests/torch_tests/test_hf_transformers.py index 1e8c3d76983..b6360ae6ca0 100644 --- a/tests/model_hub_tests/torch_tests/test_hf_transformers.py +++ b/tests/model_hub_tests/torch_tests/test_hf_transformers.py @@ -98,7 +98,7 @@ class TestTransformersModel(TestTorchConvertModel): self.image = Image.open(requests.get(url, stream=True).raw) self.cuda_available, self.gptq_postinit = None, None - def load_model(self, name, type): + def load_model_impl(self, name, type): import torch name_suffix = '' from transformers import AutoConfig diff --git a/tests/model_hub_tests/torch_tests/test_speech-transformer.py b/tests/model_hub_tests/torch_tests/test_speech-transformer.py index e08526d7a7b..a9e7013bdd2 100644 --- a/tests/model_hub_tests/torch_tests/test_speech-transformer.py +++ b/tests/model_hub_tests/torch_tests/test_speech-transformer.py @@ -24,7 +24,7 @@ class TestSpeechTransformerConvertModel(TestTorchConvertModel): checkpoint_url = "https://github.com/foamliu/Speech-Transformer/releases/download/v1.0/speech-transformer-cn.pt" subprocess.check_call(["wget", "-nv", checkpoint_url], cwd=self.repo_dir.name) - def load_model(self, model_name, model_link): + def load_model_impl(self, model_name, model_link): sys.path.append(self.repo_dir.name) from transformer.transformer import Transformer diff --git a/tests/model_hub_tests/torch_tests/test_timm.py b/tests/model_hub_tests/torch_tests/test_timm.py index 5805bfb82dd..bfea3c22902 100644 --- a/tests/model_hub_tests/torch_tests/test_timm.py +++ b/tests/model_hub_tests/torch_tests/test_timm.py @@ -10,6 +10,9 @@ from models_hub_common.constants import hf_hub_cache_dir from models_hub_common.utils import cleanup_dir, get_models_list from torch_utils import TestTorchConvertModel, process_pytest_marks +from openvino import convert_model +from torch.export import export +from packaging import version def filter_timm(timm_list: list) -> list: @@ -47,7 +50,7 @@ torch.manual_seed(0) class TestTimmConvertModel(TestTorchConvertModel): - def load_model(self, model_name, model_link): + def load_model_impl(self, model_name, model_link): m = timm.create_model(model_name, pretrained=True) cfg = timm.get_pretrained_cfg(model_name) shape = [1] + list(cfg.input_size) @@ -78,11 +81,14 @@ class TestTimmConvertModel(TestTorchConvertModel): "sequencer2d_l.in1k"]) @pytest.mark.precommit def test_convert_model_precommit(self, name, ie_device): + self.mode = "trace" self.run(name, None, ie_device) @pytest.mark.nightly + @pytest.mark.parametrize("mode", ["trace", "export"]) @pytest.mark.parametrize("name", get_all_models()) - def test_convert_model_all_models(self, name, ie_device): + def test_convert_model_all_models(self, mode, name, ie_device): + self.mode = mode self.run(name, None, ie_device) @pytest.mark.nightly diff --git a/tests/model_hub_tests/torch_tests/test_torchbench.py b/tests/model_hub_tests/torch_tests/test_torchbench.py index 04c2c96de0a..c0fba6b605e 100644 --- a/tests/model_hub_tests/torch_tests/test_torchbench.py +++ b/tests/model_hub_tests/torch_tests/test_torchbench.py @@ -30,7 +30,7 @@ class TestTorchbenchmarkConvertModel(TestTorchConvertModel): subprocess.check_call( ["git", "checkout", "850364ac2678b2363f086b7549254b6cb7df2e4d"], cwd=self.repo_dir.name) - def load_model(self, model_name, model_link): + def load_model_impl(self, model_name, model_link): subprocess.check_call([sys.executable, "install.py"] + [model_name], cwd=self.repo_dir.name) sys.path.append(self.repo_dir.name) from torchbenchmark import load_model_by_name diff --git a/tests/model_hub_tests/torch_tests/test_torchvision_models.py b/tests/model_hub_tests/torch_tests/test_torchvision_models.py index 56ee127c3c7..eb075baeec0 100644 --- a/tests/model_hub_tests/torch_tests/test_torchvision_models.py +++ b/tests/model_hub_tests/torch_tests/test_torchvision_models.py @@ -59,7 +59,7 @@ class TestTorchHubConvertModel(TestTorchConvertModel): if os.environ.get('USE_SYSTEM_CACHE', 'True') == 'False': torch.hub.set_dir(str(self.cache_dir.name)) - def load_model(self, model_name, model_link): + def load_model_impl(self, model_name, model_link): m = torch.hub.load("pytorch/vision", model_name, weights='DEFAULT', skip_validation=True) m.eval() @@ -105,10 +105,19 @@ class TestTorchHubConvertModel(TestTorchConvertModel): @pytest.mark.parametrize("model_name", ["efficientnet_b7", "raft_small", "swin_v2_s"]) @pytest.mark.precommit def test_convert_model_precommit(self, model_name, ie_device): + self.mode = "trace" self.run(model_name, None, ie_device) + @pytest.mark.parametrize("model_name", ["efficientnet_b7"]) + @pytest.mark.precommit + def test_convert_model_precommit_export(self, model_name, ie_device): + self.mode = "export" + self.run(model_name, None, ie_device) + + @pytest.mark.parametrize("mode", ["trace", "export"]) @pytest.mark.parametrize("name", process_pytest_marks(os.path.join(os.path.dirname(__file__), "torchvision_models"))) @pytest.mark.nightly - def test_convert_model_all_models(self, name, ie_device): + def test_convert_model_all_models(self, mode, name, ie_device): + self.mode = mode self.run(name, None, ie_device) diff --git a/tests/model_hub_tests/torch_tests/torch_utils.py b/tests/model_hub_tests/torch_tests/torch_utils.py index 87c9aeb043f..92bc08ee73e 100644 --- a/tests/model_hub_tests/torch_tests/torch_utils.py +++ b/tests/model_hub_tests/torch_tests/torch_utils.py @@ -45,10 +45,19 @@ def extract_unsupported_ops_from_exception(e: str) -> list: class TestTorchConvertModel(TestConvertModel): + cached_model = None def setup_class(self): torch.set_grad_enabled(False) def load_model(self, model_name, model_link): + if self.cached_model is not None and self.cached_model[0] == model_name and self.cached_model[1] == model_link: + return self.cached_model[2] + else: + res = self.load_model_impl(model_name, model_link) + self.cached_model = (model_name, model_link, res) + return res + + def load_model_impl(self, model_name, model_link): raise "load_model is not implemented" def get_inputs_info(self, model_obj): @@ -61,10 +70,40 @@ class TestTorchConvertModel(TestConvertModel): else: return flattenize_structure(inputs) + def convert_model_impl(self, model_obj): + if hasattr(self, "mode") and self.mode == "export": + from torch.fx.experimental.proxy_tensor import make_fx + from torch.export import export + from packaging import version + from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder + + graph = export(model_obj, self.example) + if version.parse(torch.__version__) >= version.parse("2.2"): + graph = graph.run_decompositions() + + try: + gm = make_fx(graph)(*self.example) + except: + gm = make_fx(graph, tracing_mode='symbolic')(*self.example) + + input_shapes = [] + input_types = [] + for input_data in self.example: + input_types.append(input_data.type()) + input_shapes.append(input_data.size()) + + decoder = TorchFXPythonDecoder(gm, gm, input_shapes=input_shapes, input_types=input_types) + ov_model = convert_model(decoder, example_input=self.example) + else: + ov_model = convert_model(model_obj, + example_input=self.example, + verbose=True + ) + return ov_model + def convert_model(self, model_obj): try: - ov_model = convert_model( - model_obj, example_input=self.example, verbose=True) + ov_model = self.convert_model_impl(model_obj) except Exception as e: report_filename = os.environ.get("OP_REPORT_FILE", None) if report_filename: diff --git a/tools/ovc/openvino/tools/ovc/convert_impl.py b/tools/ovc/openvino/tools/ovc/convert_impl.py index dbfb2c440a6..8a25b739567 100644 --- a/tools/ovc/openvino/tools/ovc/convert_impl.py +++ b/tools/ovc/openvino/tools/ovc/convert_impl.py @@ -190,12 +190,13 @@ def check_model_object(argv): return "tf" if 'torch' in sys.modules: import torch - if isinstance(model, (torch.nn.Module, torch.jit.ScriptFunction)): + if isinstance(model, (torch.nn.Module, torch.jit.ScriptFunction, torch.export.ExportedProgram)): return "pytorch" try: from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder + from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder - if isinstance(model, TorchScriptPythonDecoder): + if isinstance(model, (TorchScriptPythonDecoder, TorchFXPythonDecoder)): return "pytorch" except Exception as e: pass diff --git a/tools/ovc/openvino/tools/ovc/moc_frontend/pytorch_frontend_utils.py b/tools/ovc/openvino/tools/ovc/moc_frontend/pytorch_frontend_utils.py index 3a24e84af1a..138e9a55420 100644 --- a/tools/ovc/openvino/tools/ovc/moc_frontend/pytorch_frontend_utils.py +++ b/tools/ovc/openvino/tools/ovc/moc_frontend/pytorch_frontend_utils.py @@ -13,6 +13,8 @@ from openvino.tools.ovc.error import Error def get_pytorch_decoder(model, example_inputs, args): try: from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder + from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder + import torch except Exception as e: log.error("PyTorch frontend loading failed") raise e @@ -31,8 +33,11 @@ def get_pytorch_decoder(model, example_inputs, args): raise RuntimeError( "NNCF models produced by nncf<2.6 are not supported directly. Please upgrade nncf or export to ONNX first.") inputs = prepare_torch_inputs(example_inputs) - if not isinstance(model, TorchScriptPythonDecoder): - decoder = TorchScriptPythonDecoder(model, example_input=inputs, shared_memory=args.get("share_weights", True)) + if not isinstance(model, (TorchScriptPythonDecoder, TorchFXPythonDecoder)): + if isinstance(model, torch.export.ExportedProgram): + raise RuntimeException("Models recieved from torch.export are not yet supported by convert_model.") + else: + decoder = TorchScriptPythonDecoder(model, example_input=inputs, shared_memory=args.get("share_weights", True)) else: decoder = model args['input_model'] = decoder