natas_hw
/
openvino
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[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 <ekaterina.aidova@intel.com>

* 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 <ekaterina.aidova@intel.com>
This commit is contained in:
Maxim Vafin 2024-01-30 12:50:18 +01:00 committed by GitHub
parent 688279ab37
commit d018779add
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
86 changed files with 1283 additions and 284 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
.github/workflows/job_python_unit_tests.yml vendored
View File

@ -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: |

221
src/bindings/python/src/openvino/frontend/pytorch/fx_decoder.py
View File

@ -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:

12
src/bindings/python/src/openvino/frontend/pytorch/ts_decoder.py
View File

@ -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

5
src/bindings/python/src/openvino/frontend/pytorch/utils.py
View File

@ -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,

17
src/bindings/python/src/pyopenvino/frontend/pytorch/decoder.hpp
View File

@ -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);

2
src/bindings/python/src/pyopenvino/utils/utils.cpp
View File

@ -375,6 +375,8 @@ ov::Any py_object_to_any(const py::object& py_obj) {
return py::cast<ov::Affinity>(py_obj);
} else if (py::isinstance<ov::Tensor>(py_obj)) {
return py::cast<ov::Tensor>(py_obj);
} else if (py::isinstance<ov::Output<ov::Node>>(py_obj)) {
return py::cast<ov::Output<ov::Node>>(py_obj);
// FrontEnd Decoder
} else if (py::isinstance<ov::frontend::IDecoder>(py_obj)) {
return py::cast<std::shared_ptr<ov::frontend::IDecoder>>(py_obj);

18
src/frontends/pytorch/include/openvino/frontend/pytorch/decoder.hpp
View File

@ -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<Node> mark_node(std::shared_ptr<Node> 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;
};

31
src/frontends/pytorch/include/openvino/frontend/pytorch/node_context.hpp
View File

@ -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<Node> 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<size_t>(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<Node> 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<Output<Node>>()) {
// Case when input is constant value
return attr.as<Output<Node>>();
} else if (attr.is<type::PyNone>()) {
// 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<Node> ov_output) const;

1
src/frontends/pytorch/src/frontend.cpp
View File

@ -194,6 +194,7 @@ void FrontEnd::normalize(const std::shared_ptr<ov::Model>& model) const {
manager.register_pass<ov::frontend::pytorch::pass::PrimListUnpackReplacer>();
manager.register_pass<ov::frontend::pytorch::pass::AtenGetItemReplacer>();
manager.register_pass<ov::frontend::pytorch::pass::ListConstructReplacer>();
// TODO: remove AtenIndexToSelect when problem with dynamic input rank is gone.
manager.register_pass<ov::frontend::pytorch::pass::AtenIndexToSelect>();
manager.register_pass<ov::frontend::pytorch::pass::AtenIndexPutReplacer>();
manager.register_pass<ov::frontend::pytorch::pass::PrimListConstructPadReplacer>();

2
src/frontends/pytorch/src/node_context.cpp
View File

@ -109,7 +109,7 @@ Output<Node> NodeContext::get_tensor_from_model_or_create_input(size_t index) co
}
Output<Node> 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<int>(index));
auto input_node = input_tensor.get_node_shared_ptr();
if (std::dynamic_pointer_cast<v0::Parameter>(input_node)) {

15
src/frontends/pytorch/src/op/adaptive_poolnd.cpp
View File

@ -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

8
src/frontends/pytorch/src/op/add.cpp
View File

@ -34,8 +34,14 @@ OutputVector translate_add_common(const NodeContext& context, bool inplace) {
} else {
align_eltwise_input_types(context, lhs, rhs, true);
}
Output<Node> alpha;
if (!context.input_is_none(2)) {
auto converted_alpha = context.mark_node(std::make_shared<v1::ConvertLike>(context.get_input(2), rhs));
alpha = context.get_input(2);
} else if (context.has_attribute("alpha")) {
alpha = context.get_attribute<Output<Node>>("alpha");
}
if (alpha.get_node_shared_ptr()) {
auto converted_alpha = context.mark_node(std::make_shared<v1::ConvertLike>(alpha, rhs));
rhs = context.mark_node(std::make_shared<v1::Multiply>(converted_alpha, rhs));
}
auto add = context.mark_node(std::make_shared<v1::Add>(lhs, rhs));

21
src/frontends/pytorch/src/op/addcmul.cpp
View File

@ -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<Node>& value) {
const auto eltwise_mult = std::make_shared<v1::Multiply>(context.get_input(1), context.get_input(2));
const auto value = context.get_input(3);
const auto converted_value = std::make_shared<v1::ConvertLike>(value, context.get_input(1));
const auto scalar_mult = std::make_shared<v1::Multiply>(eltwise_mult, converted_value);
context.mark_nodes({eltwise_mult, converted_value, scalar_mult});
return {context.mark_node(std::make_shared<v1::Add>(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<Node> 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

34
src/frontends/pytorch/src/op/addmm.cpp
View File

@ -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<Node> beta, const Output<Node> 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<v0::MatMul>(m1, m2));
auto beta_converted = context.mark_node(std::make_shared<v1::ConvertLike>(beta, input));
auto alpha_converted = context.mark_node(std::make_shared<v1::ConvertLike>(alpha, mm));
auto input_beta = context.mark_node(std::make_shared<v1::Multiply>(input, beta_converted));
auto mm_alpha = context.mark_node(std::make_shared<v1::Multiply>(mm, alpha_converted));
return {context.mark_node(std::make_shared<v1::Add>(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<Node> alpha = one;
ov::Output<Node> 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<v1::ConvertLike>(beta, input));
auto alpha_converted = context.mark_node(std::make_shared<v1::ConvertLike>(alpha, mm));
auto input_beta = context.mark_node(std::make_shared<v1::Multiply>(input, beta_converted));
auto mm_alpha = context.mark_node(std::make_shared<v1::Multiply>(mm, alpha_converted));
return {context.mark_node(std::make_shared<v1::Add>(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<Node> alpha = one;
ov::Output<Node> 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

132
src/frontends/pytorch/src/op/batch_norm.cpp
View File

@ -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<Node> broadcast_const_to_channel_dim(const NodeContext& context,
auto channel_dim_exp = context.mark_node(std::make_shared<v0::Unsqueeze>(channel_dim, zero_i));
return context.mark_node(std::make_shared<v3::Broadcast>(value, channel_dim_exp));
}
OutputVector make_batch_norm(const NodeContext& context,
const Output<Node>& input,
const Output<Node>& weight,
const Output<Node>& bias,
const Output<Node>& running_mean,
const Output<Node>& running_var,
float epsilon) {
Output<Node> w = weight;
Output<Node> b = bias;
Output<Node> mean = running_mean;
Output<Node> 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<Node> rank = std::get<1>(get_shape_rank(context, input, true));
auto after_channel_dims = context.mark_node(std::make_shared<v0::Range>(two, rank, one));
auto axes = context.mark_node(std::make_shared<v0::Concat>(OutputVector{zero_1d, after_channel_dims}, 0));
if (!mean.get_node_shared_ptr()) {
mean = context.mark_node(std::make_shared<v1::ReduceMean>(input, axes, false));
}
if (!var.get_node_shared_ptr()) {
auto current_mean = context.mark_node(std::make_shared<v1::ReduceMean>(input, axes, true));
auto sub_v = context.mark_node(std::make_shared<v1::Subtract>(input, current_mean));
auto sqr_sub = context.mark_node(std::make_shared<v1::Multiply>(sub_v, sub_v));
var = context.mark_node(std::make_shared<v1::ReduceMean>(sqr_sub, axes, false));
}
return {context.mark_node(std::make_shared<v5::BatchNormInference>(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<Node> weight;
Output<Node> bias;
Output<Node> running_mean;
Output<Node> running_var;
Output<Node> current_mean;
Output<Node> 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<v3::ShapeOf>(input, element::i32));
auto rank_unsq = context.mark_node(std::make_shared<v3::ShapeOf>(input_shape, element::i32));
auto rank = context.mark_node(std::make_shared<v0::Squeeze>(rank_unsq, zero));
auto after_channel_dims = context.mark_node(std::make_shared<v0::Range>(two, rank, one));
auto axes = context.mark_node(std::make_shared<v0::Concat>(OutputVector{zero_1d, after_channel_dims}, 0));
current_mean = context.mark_node(std::make_shared<v1::ReduceMean>(input, axes, false));
auto mean = context.mark_node(std::make_shared<v1::ReduceMean>(input, axes, true));
auto sub_v = context.mark_node(std::make_shared<v1::Subtract>(input, mean));
auto sqr_sub = context.mark_node(std::make_shared<v1::Multiply>(sub_v, sub_v));
current_var = context.mark_node(std::make_shared<v1::ReduceMean>(sqr_sub, axes, false));
}
auto training = context.const_input<bool>(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<float>(6);
} else {
epsilon = context.const_input<float>(7);
}
float epsilon = context.const_input<float>(7);
// Input with index 8 is flag "cudnn_enabled" we can ignore it
return {context.mark_node(
std::make_shared<v5::BatchNormInference>(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<bool>(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<bool>(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<Node> weight;
Output<Node> 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<float>(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<Node> weight;
if (!context.input_is_none(1)) {
weight = context.get_input(1);
}
Output<Node> bias;
if (!context.input_is_none(2)) {
bias = context.get_input(2);
}
auto training = context.const_input<bool>(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<float>(5);
auto output = make_batch_norm(context, context.get_input(0), weight, bias, {}, {}, eps);
return {context.mark_node(make_list_construct(output))};
}

46
src/frontends/pytorch/src/op/div.cpp
View File

@ -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<std::string>(2);
}
OutputVector translate_div_common(const NodeContext& context,
const Output<Node>& lhs,
const Output<Node>& 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<v1::ConvertLike>(x, y));
y = context.mark_node(std::make_shared<v1::ConvertLike>(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<std::string>(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<std::string>(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<std::string>("rounding_mode");
}
return translate_div_common(context, x, y, rounding_mode, false);
};
} // namespace op

2
src/frontends/pytorch/src/op/expand.cpp
View File

@ -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<Shape>(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);

15
src/frontends/pytorch/src/op/full.cpp
View File

@ -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<element::Type>("dtype");
filled_tensor = context.mark_node(std::make_shared<v0::Convert>(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);

29
src/frontends/pytorch/src/op/gelu.cpp
View File

@ -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<std::string>(1);
}
if (approximate == "none") {
return {context.mark_node(std::make_shared<ov::op::v7::Gelu>(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<std::string>(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<std::string>("approximate");
}
return translate_gelu_common(context, approximate);
};
} // namespace op
} // namespace pytorch

2
src/frontends/pytorch/src/op/glu.cpp
View File

@ -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);

260
src/frontends/pytorch/src/op/index.cpp
View File

@ -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<Node> flatten(ov::pass::NodeRegistry& rg, const Output<Node>& 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<Node> 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<v3::ShapeOf>(value, element::i32);
const auto value_rank = rg.make<v3::ShapeOf>(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<v8::Slice>(value_shape, start, axis_node, step);
auto zero = v0::Constant::create(element::i32, {}, {0});
auto first_part_dims_length = rg.make<v1::ReduceProd>(first_part_dims, zero, true);
auto remaining_part_length = v0::Constant::create(element::i32, {1}, {-1});
output_shape = rg.make<v0::Concat>(OutputVector{first_part_dims_length, remaining_part_length}, 0);
}
return rg.make<v1::Reshape>(value, output_shape, true);
}
OutputVector index_on_list(ov::pass::NodeRegistry& rg,
const Output<Node>& data,
std::deque<Output<Node>> 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<size_t> advanced_ids;
std::vector<bool> 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<v0::Convert>(ids[i], element::u8);
auto nonzero = rg.make<v3::NonZero>(idx, element::i32);
auto input_order = v0::Constant::create(element::i32, Shape{2}, {1, 0});
auto masked_id = rg.make<v1::Transpose>(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<v8::GatherND>(data, index);
return {gather};
}
index = rg.make<v0::Convert>(index, element::i32);
auto dim = v0::Constant::create(element::i32, Shape{}, {advanced_ids[0]});
auto gather = rg.make<v8::Gather>(data, index, dim);
return {gather};
}
auto adv_idx_count = advanced_ids.size();
auto input_shape = rg.make<v3::ShapeOf>(data, element::i32);
auto zero = v0::Constant::create(element::i32, Shape{}, {0});
auto input_dims = rg.make<v1::Split>(input_shape, zero, rank);
std::vector<size_t> 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<size_t> 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<v1::Transpose>(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<v0::Convert>(cum_adv_index, element::i32);
auto multiplier = input_dims->output(advanced_ids.back());
for (int i = static_cast<int>(adv_idx_count) - 2; i > -1; i--) {
auto input_id = advanced_ids[i];
auto m_idx = rg.make<v0::Convert>(masked_indicies[input_id], element::i32);
auto adv_index = rg.make<v1::Multiply>(m_idx, multiplier);
cum_adv_index = rg.make<v1::Add>(cum_adv_index, adv_index);
multiplier = rg.make<v1::Multiply>(multiplier, input_dims->output(input_id));
}
std::shared_ptr<Node> gather = rg.make<v8::Gather>(flatten_input, cum_adv_index, zero);
OutputVector concat_dims;
// check if all advanced indices are consecutive.
std::vector<size_t> consequence_dims;
auto cum_adv_index_shape_tensor = rg.make<v3::ShapeOf>(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<v0::Concat>(folded_adv_idx_shape_vector, 0);
gather = rg.make<v1::Reshape>(gather, folded_adv_idx_shape, false);
std::vector<size_t> 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<v1::Transpose>(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<v0::Concat>(concat_dims, 0);
gather = rg.make<v1::Reshape>(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<type::List>()) {
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<element::Type>()) {
index_ov_type = index_dtype.as<element::Type>();
}
}
if (index_ov_type == element::boolean || index_ov_type == element::u8) {
auto nonzero = context.mark_node(std::make_shared<v3::NonZero>(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<v1::Transpose>(nonzero, input_order));
auto gather = context.mark_node(std::make_shared<v8::GatherND>(x, masked_id));
return {gather};
}
if (index_ov_type != element::i32) {
indices = context.mark_node(std::make_shared<ov::op::v0::Convert>(indices, element::i32));
}
auto dim = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0}));
return {context.mark_node(std::make_shared<v8::Gather>(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<Output<Node>> list_elems;
for (size_t i = 1; i < context.get_input_size(); i++) {
auto index = context.get_input(static_cast<int>(i));
if (index.get_element_type() == element::i64) {
auto converted = context.mark_node(std::make_shared<ov::op::v0::Convert>(index, element::i32));
list_elems.push_back(converted);
} else {
list_elems.push_back(index);
Output<Node> index;
if (!context.input_is_none(i)) {
index = context.get_input(static_cast<int>(i));
}
list_elems.push_back(index);
}
auto concat =
context.mark_node(std::make_shared<ov::op::v0::Concat>(OutputVector(list_elems.begin(), list_elems.end()), 0));
auto gather = std::make_shared<ov::op::v8::GatherND>(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

32
src/frontends/pytorch/src/op/leaky_relu.cpp Normal file
View File

@ -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<Node> 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<ov::op::v1::ConvertLike>(negative_slope, x));
} else {
negative_slope = context.get_input(1);
}
return {context.mark_node(std::make_shared<v0::PRelu>(x, negative_slope))};
};
} // namespace op
} // namespace pytorch
} // namespace frontend
} // namespace ov

26
src/frontends/pytorch/src/op/mean.cpp
View File

@ -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<Node> 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<bool>(2);
}
if (!context.input_is_none(3)) {
x = apply_dtype(context, 3, x);
}
}
auto mean = context.mark_node(std::make_shared<ov::op::v1::ReduceMean>(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

52
src/frontends/pytorch/src/op/pad.cpp
View File

@ -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<std::vector<int64_t>>(1);
std::string mode = "constant";
namespace {
OutputVector translate_pad_common(const NodeContext& context,
const Output<Node>& data,
const std::vector<int64_t>& paddings,
const Output<Node>& pad_value,
const std::string& mode = "constant") {
Output<Node> shape;
Output<Node> 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<v3::Broadcast>(zero, pads_diff));
auto pads_begins = context.mark_node(std::make_shared<v0::Concat>(NodeVector{pads_remaining, pads_begin_short}, 0));
auto pads_ends = context.mark_node(std::make_shared<v0::Concat>(NodeVector{pads_remaining, pads_end_short}, 0));
if (!context.input_is_none(2)) {
mode = context.const_input<std::string>(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<v1::ConvertLike>(pad_value, data));
const std::map<std::string, PadMode> pt_to_ov_pad{
{"constant", PadMode::CONSTANT},
{"reflect", PadMode::REFLECT},
{"replicate", PadMode::EDGE},
};
Output<Node> 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<v1::ConvertLike>(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<v1::Pad>(data, pads_begins, pads_ends, pad_value, ov_mode->second))};
return {context.mark_node(std::make_shared<v1::Pad>(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<std::vector<int64_t>>(1);
std::string mode = "constant";
if (!context.input_is_none(2)) {
mode = context.const_input<std::string>(2);
}
Output<Node> 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<std::vector<int64_t>>(1);
auto pad_value = context.get_input(2);
return translate_pad_common(context, data, paddings, pad_value);
}
} // namespace op

34
src/frontends/pytorch/src/op/scalar_tensor.cpp Normal file
View File

@ -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<element::Type>("dtype");
data = context.mark_node(std::make_shared<v0::Convert>(context.get_input(0), dtype));
}
// layout and device can be ignored
return {data};
}
} // namespace op
} // namespace pytorch
} // namespace frontend
} // namespace ov

40
src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp
View File

@ -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<bool>("scale");
} else if (!context.input_is_none(4)) {
causal = context.const_input<bool>(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<v1::ConvertLike>(zero, query)));
}
if (context.has_attribute("scale")) {
const auto scale = context.get_input("scale");
inputs.push_back(context.mark_node(std::make_shared<v1::ConvertLike>(scale, query)));
}
auto sdpa = context.mark_node(std::make_shared<v13::ScaledDotProductAttention>(inputs, causal));
return {context.mark_node(make_list_construct({sdpa}))};
};
} // namespace op

68
src/frontends/pytorch/src/op/slice_scatter.cpp Normal file
View File

@ -0,0 +1,68 @@
// Copyright (C) 2018-2023 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include <climits>
#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<ov::Node> 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<v0::Unsqueeze>(dim, axis_0));
}
} else {
dim = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {1}));
}
ov::Output<ov::Node> 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<v0::Unsqueeze>(start, axis_0));
}
} else {
start = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0}));
}
ov::Output<ov::Node> 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<v0::Unsqueeze>(end, axis_0));
}
} else {
end = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {INT_MAX}));
}
ov::Output<ov::Node> 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<v0::Unsqueeze>(step, axis_0));
}
} else {
step = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {1}));
}
return {context.mark_node(std::make_shared<SliceAssign>(input, src, start, end, step, dim))};
};
} // namespace op
} // namespace pytorch
} // namespace frontend
} // namespace ov

16
src/frontends/pytorch/src/op/split.cpp
View File

@ -4,10 +4,11 @@
#include "openvino/op/split.hpp"
#include <climits>
//#include <climits>
#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<int>(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<v1::VariadicSplit>(data, dim, split_lengths);
return {context.mark_node(make_list_construct(split->outputs()))};
}
} // namespace op
} // namespace pytorch
} // namespace frontend

16
src/frontends/pytorch/src/op/sub.cpp
View File

@ -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<v1::ConvertLike>(x, y));
y = context.mark_node(std::make_shared<v1::ConvertLike>(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<Output<Node>>("alpha");
auto casted_alpha = context.mark_node(std::make_shared<v1::ConvertLike>(alpha, y));
y = context.mark_node(std::make_shared<v1::Multiply>(casted_alpha, y));
}
return {context.mark_node(std::make_shared<v1::Subtract>(x, y))};
};
} // namespace op
} // namespace pytorch
} // namespace frontend

10
src/frontends/pytorch/src/op/to.cpp
View File

@ -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<element::Type>("dtype");
data = context.mark_node(std::make_shared<v0::Convert>(context.get_input(0), dtype));
}
return {data};
}
} // namespace op
} // namespace pytorch
} // namespace frontend

86
src/frontends/pytorch/src/op_table.cpp
View File

@ -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<std::string, CreatorFunction> get_supported_ops_ts() {
const std::map<std::string, CreatorFunction> get_supported_ops_fx() {
return {
{"aten._adaptive_avg_pool2d.default", op::translate_1to1_match_2_inputs<opset10::AdaptiveAvgPool>},
{"aten.abs.default", op::translate_1to1_match_1_inputs<opset10::Abs>},
{"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<opset10::MatMul>},
{"aten.cat.default", op::translate_cat_fx},
{"aten.ceil.default", op::translate_1to1_match_1_inputs<opset10::Ceiling>},
{"aten.clamp.default", op::translate_clamp},
{"aten.clamp_min.default", op::translate_1to1_match_2_inputs<opset10::Maximum>},
{"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<opset10::Cos>},
{"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<opset10::Equal>},
{"aten.eq.Tensor", op::translate_1to1_match_2_inputs_align_types<opset10::Equal>},
{"aten.exp.default", op::translate_1to1_match_1_inputs<opset10::Exp>},
{"aten.expand.default", op::translate_expand_fx},
{"aten.full.default", op::translate_full},
{"aten.floor.default", op::translate_1to1_match_1_inputs<opset10::Floor>},
{"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<opset10::Greater>},
{"aten.hardsigmoid.default", op::translate_1to1_match_1_inputs<opset10::HSigmoid>},
{"aten.hardswish.default", op::translate_1to1_match_1_inputs<opset10::HSwish>},
{"aten.hardswish_.default", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::HSwish>>},
{"aten.hardtanh.default", op::translate_hardtanh},
{"aten.hardtanh_.default", op::inplace_op<op::translate_hardtanh>},
{"aten.index.Tensor", op::translate_index_fx},
{"aten.leaky_relu_.default", op::inplace_op<op::translate_1to1_match_2_inputs<opset10::PRelu>>},
{"aten.leaky_relu.default", op::translate_leaky_relu_fx},
{"aten.leaky_relu_.default", op::inplace_op<op::translate_leaky_relu_fx>},
{"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<opset10::Less>},
{"aten.lt.Tensor", op::translate_1to1_match_2_inputs_align_types<opset10::Less>},
{"aten.masked_fill_.Scalar", op::inplace_op<op::translate_masked_fill>},
{"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<opset10::MatMul>},
{"aten.mul.Tensor", op::translate_1to1_match_2_inputs_align_types<opset10::Multiply>},
{"aten.mul.Scalar", op::translate_1to1_match_2_inputs_align_types<opset10::Multiply>},
{"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<opset10::NotEqual>},
{"aten.neg.default", op::translate_neg},
{"aten.new_ones.default", op::translate_new_ones},
{"aten.permute.default", op::translate_1to1_match_2_inputs<opset10::Transpose>},
{"aten.pow.Tensor_Scalar", op::translate_pow},
{"aten.relu.default", op::translate_1to1_match_1_inputs<opset10::Relu>},
{"aten.relu_.default", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::Relu>>},
{"aten.repeat.default", op::translate_1to1_match_2_inputs<opset10::Tile>},
{"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<opset10::Sigmoid>},
{"aten.silu.default", op::translate_1to1_match_1_inputs<opset10::Swish>},
{"aten.silu_.default", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::Swish>>},
{"aten.sin.default", op::translate_1to1_match_1_inputs<opset10::Sin>},
{"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<opset10::Tanh>},
{"aten.unfold.default", op::translate_unfold},
{"aten.transpose.int", op::translate_transpose},
{"aten.unsqueeze.default", op::translate_1to1_match_2_inputs<opset10::Unsqueeze>},
{"aten.upsample_nearest2d.default", op::translate_upsample_nearest2d},

19
src/frontends/pytorch/src/translate_session.cpp
View File

@ -74,10 +74,6 @@ std::shared_ptr<Model> TranslateSession::convert_pytorch_model(
std::shared_ptr<TorchDecoder> pytorch_model,
const TensorMap& external_tensor_map,
const std::shared_ptr<pytorch::InputModel>& 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<Model> resulting_model; // define here to make a conversion in a nested scope
{
auto parameters = std::make_shared<ParameterVector>();
@ -128,18 +124,6 @@ std::shared_ptr<Model> 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);
}

2
src/frontends/pytorch/src/utils.cpp
View File

@ -346,7 +346,7 @@ std::shared_ptr<ov::op::util::FrameworkNode> cast_fw_node(std::shared_ptr<Node>
return fw_node;
}
std::shared_ptr<ov::op::util::FrameworkNode> make_list_construct(const ov::OutputVector& inputs) {
std::shared_ptr<ov::Node> 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");

15
src/frontends/pytorch/src/utils.hpp
View File

@ -64,7 +64,7 @@ OutputVector make_framework_node(const NodeContext& context, const std::string&
std::shared_ptr<op::util::FrameworkNode> cast_fw_node(std::shared_ptr<Node> node, const std::string& type);
std::shared_ptr<op::util::FrameworkNode> make_list_construct(const ov::OutputVector& inputs);
std::shared_ptr<Node> make_list_construct(const ov::OutputVector& inputs);
bool is_none_node(const Output<Node>& 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:

1
tests/layer_tests/pytest.ini
View File

@ -5,4 +5,5 @@ markers =
precommit_tf_fe
precommit_ts_backend
precommit_fx_backend
precommit_torch_export
timeout

64
tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py
View File

@ -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)

3
tests/layer_tests/pytorch_tests/test_adaptive_avg_pool.py
View File

@ -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):

3
tests/layer_tests/pytorch_tests/test_adaptive_max_pool.py
View File

@ -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',

1
tests/layer_tests/pytorch_tests/test_addcmul.py
View File

@ -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)

2
tests/layer_tests/pytorch_tests/test_addmm.py
View File

@ -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)

4
tests/layer_tests/pytorch_tests/test_and.py
View File

@ -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(

3
tests/layer_tests/pytorch_tests/test_as_strided.py
View File

@ -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)

1
tests/layer_tests/pytorch_tests/test_batch_norm.py
View File

@ -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)

1
tests/layer_tests/pytorch_tests/test_broadcast_tensors.py
View File

@ -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]])

1
tests/layer_tests/pytorch_tests/test_clone.py
View File

@ -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)

3
tests/layer_tests/pytorch_tests/test_conv_transposend.py
View File

@ -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)

2
tests/layer_tests/pytorch_tests/test_convolution.py
View File

@ -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)

2
tests/layer_tests/pytorch_tests/test_device.py
View File

@ -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),

2
tests/layer_tests/pytorch_tests/test_div.py
View File

@ -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'),

1
tests/layer_tests/pytorch_tests/test_embedding.py
View File

@ -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):

1
tests/layer_tests/pytorch_tests/test_eq.py
View File

@ -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]

3
tests/layer_tests/pytorch_tests/test_expand.py
View File

@ -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)

1
tests/layer_tests/pytorch_tests/test_flatten.py
View File

@ -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)

2
tests/layer_tests/pytorch_tests/test_floor_divide.py
View File

@ -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:

1
tests/layer_tests/pytorch_tests/test_gelu.py
View File

@ -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)

1
tests/layer_tests/pytorch_tests/test_glu.py
View File

@ -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)

1
tests/layer_tests/pytorch_tests/test_group_norm.py
View File

@ -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)

27
tests/layer_tests/pytorch_tests/test_index.py
View File

@ -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)

3
tests/layer_tests/pytorch_tests/test_mm.py
View File

@ -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)

2
tests/layer_tests/pytorch_tests/test_mul.py
View File

@ -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

2
tests/layer_tests/pytorch_tests/test_permute.py
View File

@ -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})

1
tests/layer_tests/pytorch_tests/test_reshape.py
View File

@ -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)

1
tests/layer_tests/pytorch_tests/test_reshape_as.py
View File

@ -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):

2
tests/layer_tests/pytorch_tests/test_resolve_conj_neg.py
View File

@ -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):

22
tests/layer_tests/pytorch_tests/test_scaled_dot_product_attention.py
View File

@ -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})

1
tests/layer_tests/pytorch_tests/test_select.py
View File

@ -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)

1
tests/layer_tests/pytorch_tests/test_silu.py
View File

@ -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)

2
tests/layer_tests/pytorch_tests/test_softmax.py
View File

@ -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:

1
tests/layer_tests/pytorch_tests/test_sub.py
View File

@ -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

1
tests/layer_tests/pytorch_tests/test_tensor_split.py
View File

@ -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)

1
tests/layer_tests/pytorch_tests/test_type_as.py
View File

@ -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})

2
tests/layer_tests/pytorch_tests/test_unflatten.py
View File

@ -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})

1
tests/layer_tests/pytorch_tests/test_unfold.py
View File

@ -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),

1
tests/model_hub_tests/torch_tests/requirements.txt
View File

@ -7,6 +7,7 @@ datasets
facexlib
numpy
optimum
packaging
pandas
protobuf
pyctcdecode

2
tests/model_hub_tests/torch_tests/test_aliked.py
View File

@ -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

2
tests/model_hub_tests/torch_tests/test_detectron2.py
View File

@ -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

2
tests/model_hub_tests/torch_tests/test_edsr.py
View File

@ -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)

2
tests/model_hub_tests/torch_tests/test_gfpgan.py
View File

@ -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

2
tests/model_hub_tests/torch_tests/test_hf_transformers.py
View File

@ -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

2
tests/model_hub_tests/torch_tests/test_speech-transformer.py
View File

@ -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

10
tests/model_hub_tests/torch_tests/test_timm.py
View File

@ -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

2
tests/model_hub_tests/torch_tests/test_torchbench.py
View File

@ -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

13
tests/model_hub_tests/torch_tests/test_torchvision_models.py
View File

@ -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)

43
tests/model_hub_tests/torch_tests/torch_utils.py
View File

@ -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:

5
tools/ovc/openvino/tools/ovc/convert_impl.py
View File

@ -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

9
tools/ovc/openvino/tools/ovc/moc_frontend/pytorch_frontend_utils.py
View File

@ -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