diff --git a/src/frontends/pytorch/src/op/norm.cpp b/src/frontends/pytorch/src/op/norm.cpp index 34a0bdd01c4..a743816ee8d 100644 --- a/src/frontends/pytorch/src/op/norm.cpp +++ b/src/frontends/pytorch/src/op/norm.cpp @@ -5,12 +5,23 @@ #include "openvino/frontend/pytorch/node_context.hpp" #include "openvino/op/abs.hpp" #include "openvino/op/constant.hpp" +#include "openvino/op/convert.hpp" +#include "openvino/op/convert_like.hpp" +#include "openvino/op/gather.hpp" +#include "openvino/op/multiply.hpp" +#include "openvino/op/not_equal.hpp" #include "openvino/op/power.hpp" +#include "openvino/op/range.hpp" #include "openvino/op/reduce_l1.hpp" #include "openvino/op/reduce_l2.hpp" #include "openvino/op/reduce_max.hpp" #include "openvino/op/reduce_min.hpp" #include "openvino/op/reduce_sum.hpp" +#include "openvino/op/reshape.hpp" +#include "openvino/op/shape_of.hpp" +#include "openvino/op/sqrt.hpp" +#include "openvino/op/squeeze.hpp" +#include "pt_framework_node.hpp" #include "utils.hpp" namespace ov { @@ -20,13 +31,12 @@ namespace op { using namespace ov::op; -OutputVector translate_norm(const NodeContext& context) { - num_inputs_check(context, 4, 4); - auto input_tensor = context.get_input(0); - auto p = context.const_input(1); - auto dim = context.get_input(2); - auto keep_dim = context.const_input(3); - +namespace { +Output norm_vector(const NodeContext& context, + Output input_tensor, + Output dim, + float p, + bool keep_dim) { Output res; if (p == 1) { res = context.mark_node(std::make_shared(input_tensor, dim, keep_dim)); @@ -38,17 +48,253 @@ OutputVector translate_norm(const NodeContext& context) { } else if (p == -std::numeric_limits::infinity()) { auto abs = context.mark_node(std::make_shared(input_tensor)); res = context.mark_node(std::make_shared(abs, dim, keep_dim)); + } else if (p == 0) { + auto input_rank = input_tensor.get_partial_shape().rank(); + FRONT_END_OP_CONVERSION_CHECK(input_rank.is_dynamic() || input_rank.get_length() == 1, + "ord=0 supported only for vector norm"); + auto zero = context.mark_node(v0::Constant::create(element::f32, Shape{}, {0})); + zero = context.mark_node(std::make_shared(zero, input_tensor)); + auto cond = context.mark_node(std::make_shared(input_tensor, zero)); + cond = context.mark_node(std::make_shared(cond, input_tensor)); + res = context.mark_node(std::make_shared(cond, dim, keep_dim)); } else { - auto const_p = context.mark_node(v0::Constant::create(element::f64, Shape{1}, {p})); - auto const_p_inv = context.mark_node(v0::Constant::create(element::f64, Shape{1}, {1.0 / p})); + auto const_p = context.mark_node(v0::Constant::create(element::f32, Shape{1}, {p})); + const_p = context.mark_node(std::make_shared(const_p, input_tensor)); + auto const_p_inv = context.mark_node(v0::Constant::create(element::f32, Shape{1}, {1.0 / p})); + const_p_inv = context.mark_node(std::make_shared(const_p_inv, input_tensor)); auto abs = context.mark_node(std::make_shared(input_tensor)); auto pow = context.mark_node(std::make_shared(abs, const_p)); auto sum = context.mark_node(std::make_shared(pow, dim, keep_dim)); res = context.mark_node(std::make_shared(sum, const_p_inv)); } + return res; +}; + +Output norm_matrix(const NodeContext& context, + Output input_tensor, + Output dim, + float p, + bool keep_dim) { + Output res; + auto one = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1})); + auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); + auto first_dim = context.mark_node(std::make_shared(dim, zero, zero)); + auto second_dim = context.mark_node(std::make_shared(dim, one, zero)); + if (p == 1) { + auto abs = context.mark_node(std::make_shared(input_tensor)); + auto sum = context.mark_node(std::make_shared(abs, first_dim, true)); + res = context.mark_node(std::make_shared(sum, second_dim, keep_dim)); + } else if (p == std::numeric_limits::infinity()) { + auto abs = context.mark_node(std::make_shared(input_tensor)); + auto sum = context.mark_node(std::make_shared(abs, second_dim, true)); + res = context.mark_node(std::make_shared(sum, first_dim, keep_dim)); + } else if (p == -std::numeric_limits::infinity()) { + auto abs = context.mark_node(std::make_shared(input_tensor)); + auto sum = context.mark_node(std::make_shared(abs, second_dim, true)); + res = context.mark_node(std::make_shared(sum, first_dim, keep_dim)); + } else if (p == -1) { + auto abs = context.mark_node(std::make_shared(input_tensor)); + auto sum = context.mark_node(std::make_shared(abs, first_dim, true)); + res = context.mark_node(std::make_shared(sum, second_dim, keep_dim)); + } else { + FRONT_END_OP_CONVERSION_CHECK(false, "Unsupported ord ", p, " for matrix norm"); + } + + return res; +}; + +Output frobenius_norm(const NodeContext& context, Output x, Output dim, bool keep_dim) { + auto sqr = context.mark_node(std::make_shared(x, x)); + auto sumsqr = context.mark_node(std::make_shared(sqr, dim, keep_dim)); + return context.mark_node(std::make_shared(sumsqr)); +} +}; // namespace + +OutputVector translate_norm(const NodeContext& context) { + num_inputs_check(context, 4, 6); + auto input_tensor = context.get_input(0); + auto p_node_type = context.get_input_type(1); + Output dim; + if (context.input_is_none(2)) { + auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); + auto one = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1})); + auto input_shape = context.mark_node(std::make_shared(input_tensor, element::i32)); + auto rank = context.mark_node(std::make_shared(input_shape, element::i32)); + rank = context.mark_node(std::make_shared(rank, zero)); + dim = context.mark_node(std::make_shared(zero, rank, one)); + } else { + dim = context.get_input(2); + } + auto keep_dim = context.const_input(3); + if (!context.input_is_none(4)) { + input_tensor = apply_dtype(context, 4, input_tensor); + } + Output res; + if (p_node_type.is()) { + auto p_str = context.const_input(1); + if (p_str == "fro") { + res = frobenius_norm(context, input_tensor, dim, keep_dim); + } else { + FRONT_END_OP_CONVERSION_CHECK(false, "Umsupported ord ", p_str); + } + } else { + auto p = context.const_input(1); + res = norm_vector(context, input_tensor, dim, p, keep_dim); + } + // output tensor + if (!context.input_is_none(5)) { + context.mutate_input(5, res); + } return {res}; }; +OutputVector translate_linalg_vector_norm(const NodeContext& context) { + // aten::linalg_vector_norm(Tensor self, Scalar ord=2, int[1]? dim=None, bool keepdim=False, *, ScalarType? + // dtype=None) -> Tensor + // aten::linalg_vector_norm.out(Tensor self, Scalar ord=2, int[1]? dim=None, bool + // keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!): + num_inputs_check(context, 5, 6); + auto x = context.get_input(0); + // ord defines the vector norm that is computed. + auto ord = context.const_input(1); + bool keep_dim = context.const_input(3); + Output dim; + Output result; + // If dim= None, x will be flattened before the norm is computed. + if (context.input_is_none(2)) { + keep_dim = false; + auto minus_one = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-1})); + x = context.mark_node(std::make_shared(x, minus_one, false)); + dim = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0})); + } else { + dim = context.get_input(2); + } + // dtype may be used to perform the computation in a more precise dtype. It is semantically equivalent to calling + // linalg.vector_norm(x.to(dtype)) + if (!context.input_is_none(4)) { + x = apply_dtype(context, 4, x); + } + result = norm_vector(context, x, dim, ord, keep_dim); + // output tensor + if (!context.input_is_none(5)) { + context.mutate_input(5, result); + } + return {result}; +}; + +OutputVector translate_linalg_matrix_norm(const NodeContext& context) { + // aten::linalg_matrix_norm.out(Tensor self, Scalar ord, int[] dim=[-2, -1], bool keepdim=False, *, ScalarType? + // dtype=None, Tensor(a!) out) -> Tensor(a!) aten::linalg_matrix_norm(Tensor self, Scalar ord, int[] dim=[-2, -1], + // bool keepdim=False, *, ScalarType? dtype=None) aten::linalg_matrix_norm.str_ord(Tensor self, str ord="fro", int[] + // dim=[-2, -1], bool keepdim=False, *, ScalarType? dtype=None) + num_inputs_check(context, 5, 6); + auto x = context.get_input(0); + // ord defines the vector norm that is computed can be string or number + auto ord_type = context.get_input_type(1); + auto dim = context.get_input(2); + bool keep_dim = context.const_input(3); + Output result; + + // dtype may be used to perform the computation in a more precise dtype. It is semantically equivalent to calling + // linalg.mtrix_norm(x.to(dtype)) + if (!context.input_is_none(4)) { + x = apply_dtype(context, 4, x); + } + if (ord_type.is()) { + auto p_str = context.const_input(1); + if (p_str == "fro") { + result = frobenius_norm(context, x, dim, keep_dim); + } else { + FRONT_END_OP_CONVERSION_CHECK(false, "Umsupported ord ", p_str); + } + } else { + auto p = context.const_input(1); + result = norm_matrix(context, x, dim, p, keep_dim); + } + // output tensor + if (!context.input_is_none(5)) { + context.mutate_input(5, result); + } + return {result}; +}; + +OutputVector translate_linalg_norm(const NodeContext& context) { + // aten::linalg_norm(Tensor self, Scalar? ord=None, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) + // aten::linalg_norm.ord_str(Tensor self, str ord, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) + // aten::linalg_norm.ord_str_out(Tensor self, str ord, int[1]? dim=None, bool keepdim=False, *, ScalarType? + // dtype=None, Tensor(a!) out) -> Tensor(a!) + num_inputs_check(context, 5, 6); + auto x = context.get_input(0); + bool keep_dim = context.const_input(3); + Output result; + Output dim; + // dtype may be used to perform the computation in a more precise dtype. It is semantically equivalent to calling + // linalg.norm(x.to(dtype)) + if (!context.input_is_none(4)) { + x = apply_dtype(context, 4, x); + } + // If dim= None and ord= None, A will be flattened to 1D and the 2-norm of the resulting vector will be computed. + if (context.input_is_none(2) && context.input_is_none(1)) { + auto minus_one = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-1})); + x = context.mark_node(std::make_shared(x, minus_one, false)); + dim = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0})); + result = norm_vector(context, x, dim, 2, false); + } else { + // If dim=None apply for all dimesions + if (context.input_is_none(2)) { + auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); + auto one = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1})); + auto input_shape = context.mark_node(std::make_shared(x, element::i32)); + auto rank = context.mark_node(std::make_shared(input_shape, element::i32)); + rank = context.mark_node(std::make_shared(rank, zero)); + dim = context.mark_node(std::make_shared(zero, rank, one)); + } else { + dim = context.get_input(2); + } + // default norm for matrix is frobenius norm, for vector - L2, for other ranks are not detrmined + if (context.input_is_none(1)) { + auto input_rank = x.get_partial_shape().rank(); + if (input_rank.is_static() && input_rank.get_length() == 2) { + result = frobenius_norm(context, x, dim, keep_dim); + } else if (input_rank.is_static() && input_rank.get_length() == 1) { + result = norm_vector(context, x, dim, 2, keep_dim); + } else { + FRONT_END_OP_CONVERSION_CHECK(false, + "linalg norm for tensor rank > 2 without ord specification unsupported"); + } + } else { + // ord defines the norm that is computed can be string or number + auto ord_type = context.get_input_type(1); + if (ord_type.is()) { + auto p_str = context.const_input(1); + if (p_str == "fro") { + result = frobenius_norm(context, x, dim, keep_dim); + } else { + FRONT_END_OP_CONVERSION_CHECK(false, "Umsupported ord ", p_str); + } + } else { + auto p = context.const_input(1); + if (!context.input_is_none(2)) { + auto const_dim = context.const_input>(2); + if (const_dim.size() == 2) { + result = norm_matrix(context, x, dim, p, keep_dim); + } else { + result = norm_vector(context, x, dim, p, keep_dim); + } + } else { + result = norm_vector(context, x, dim, p, keep_dim); + } + } + } + } + + // output tensor + if (!context.input_is_none(5)) { + context.mutate_input(5, result); + } + return {result}; +}; + } // namespace op } // namespace pytorch } // namespace frontend diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp index c42024fa36f..80e1db8b354 100644 --- a/src/frontends/pytorch/src/op_table.cpp +++ b/src/frontends/pytorch/src/op_table.cpp @@ -62,6 +62,9 @@ OP_CONVERTER(translate_instance_norm); OP_CONVERTER(translate_int); OP_CONVERTER(translate_layer_norm); OP_CONVERTER(translate_len); +OP_CONVERTER(translate_linalg_norm); +OP_CONVERTER(translate_linalg_matrix_norm); +OP_CONVERTER(translate_linalg_vector_norm); OP_CONVERTER(translate_linear); OP_CONVERTER(translate_list_construct); OP_CONVERTER(translate_log); @@ -235,6 +238,9 @@ const std::map get_supported_ops() { {"aten::leaky_relu", op::translate_1to1_match_2_inputs}, {"aten::leaky_relu_", op::inplace_op>}, {"aten::len", op::translate_len}, + {"aten::linalg_norm", op::translate_linalg_norm}, + {"aten::linalg_matrix_norm", op::translate_linalg_matrix_norm}, + {"aten::linalg_vector_norm", op::translate_linalg_vector_norm}, {"aten::linear", op::translate_linear}, {"aten::log", op::translate_log}, {"aten::log_", op::inplace_op}, diff --git a/tests/layer_tests/pytorch_tests/test_norm.py b/tests/layer_tests/pytorch_tests/test_norm.py index 7cea64f4825..7005ed4af63 100644 --- a/tests/layer_tests/pytorch_tests/test_norm.py +++ b/tests/layer_tests/pytorch_tests/test_norm.py @@ -14,7 +14,7 @@ from pytorch_layer_test_class import PytorchLayerTest class TestNorm(PytorchLayerTest): def _prepare_input(self): - return (np.random.randn(2, 3, 4, 5),) + return (np.random.randn(1, 2, 3).astype(np.float32),) def create_model(self, p, dim, keepdim): class aten_norm(torch.nn.Module): @@ -37,3 +37,216 @@ class TestNorm(PytorchLayerTest): def test_norm(self, ie_device, precision, ir_version, p, dim, keepdim): self._test(*self.create_model(p, dim, keepdim), ie_device, precision, ir_version) + + +class TestLinalgVectorNorm(PytorchLayerTest): + + def _prepare_input(self, out=False, out_dtype=None): + if not out: + return (np.random.randn(1, 2, 3).astype(np.float32),) + x = np.random.randn(1, 2, 3).astype(np.float32) + y = np.random.randn(1, 2, 3).astype( + out_dtype if out_dtype is not None else np.float32) + return (x, y) + + def create_model(self, p, dim, keepdim, dtype_str, out, out_as_dtype): + dtypes = { + "float32": torch.float32, + "float64": torch.float64 + } + dtype = dtypes.get(dtype_str) + + class aten_linalg_vector_norm(torch.nn.Module): + + def __init__(self, p, dim, keepdim, dtype, out, out_as_dtype) -> None: + super().__init__() + self.ord = p + self.dim = dim + self.keepdim = keepdim + self.dtype = dtype + if self.dtype is not None: + self.forward = self.forward_dtype + if out: + self.forward = self.forward_out + if out_as_dtype: + self.forward = self.forward_prim_dtype + + def forward(self, x): + return torch.linalg.vector_norm( + x, ord=self.ord, dim=self.dim, keepdim=self.keepdim + ) + + def forward_dtype(self, x): + return torch.linalg.vector_norm( + x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=self.dtype + ) + + def forward_prim_dtype(self, x, y): + return torch.linalg.vector_norm( + x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=y.dtype + ) + + def forward_out(self, x, y): + return y, torch.linalg.vector_norm( + x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, out=y + ) + + ref_net = None + + return aten_linalg_vector_norm(p, dim, keepdim, dtype, out, out_as_dtype), ref_net, "aten::linalg_vector_norm" + + @pytest.mark.nightly + @pytest.mark.precommit + @pytest.mark.parametrize('p', [-2, -1, 0, 1, 2, 2.5, float('inf'), float('-inf')]) + @pytest.mark.parametrize('dim', [0, [0, 1], None]) + @pytest.mark.parametrize('keepdim', [True, False]) + @pytest.mark.parametrize("dtype", ["float32", "float64", None]) + @pytest.mark.parametrize("out", [True, False]) + @pytest.mark.parametrize("prim_dtype", [True, False]) + def test_linalg_vector_norm(self, p, dim, keepdim, dtype, out, prim_dtype, ie_device, precision, ir_version): + self._test(*self.create_model(p, dim, keepdim, dtype, out, prim_dtype), + ie_device, precision, ir_version, + kwargs_to_prepare_input={"out": out or prim_dtype, "out_dtype": dtype if prim_dtype else None}) + + +class TestLinalgMatrixNorm(PytorchLayerTest): + + def _prepare_input(self, out=False, out_dtype=None): + if not out: + return (np.random.randn(3, 3).astype(np.float32),) + x = np.random.randn(1, 3, 3).astype(np.float32) + y = np.random.randn(1, 3, 3).astype( + out_dtype if out_dtype is not None else np.float32) + return (x, y) + + def create_model(self, p, dim, keepdim, dtype_str, out, out_as_dtype): + dtypes = { + "float32": torch.float32, + "float64": torch.float64 + } + dtype = dtypes.get(dtype_str) + + class aten_linalg_matrix_norm(torch.nn.Module): + + def __init__(self, p, dim, keepdim, dtype, out, out_as_dtype) -> None: + super().__init__() + self.ord = p + self.dim = dim + self.keepdim = keepdim + self.dtype = dtype + if self.dtype is not None: + self.forward = self.forward_dtype + if out: + self.forward = self.forward_out + if out_as_dtype: + self.forward = self.forward_prim_dtype + + def forward(self, x): + return torch.linalg.matrix_norm( + x, ord=self.ord, dim=self.dim, keepdim=self.keepdim + ) + + def forward_dtype(self, x): + return torch.linalg.matrix_norm( + x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=self.dtype + ) + + def forward_prim_dtype(self, x, y): + return torch.linalg.matrix_norm( + x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=y.dtype + ) + + def forward_out(self, x, y): + return y, torch.linalg.matrix_norm( + x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, out=y + ) + + ref_net = None + + return aten_linalg_matrix_norm(p, dim, keepdim, dtype, out, out_as_dtype), ref_net, "aten::linalg_matrix_norm" + + @pytest.mark.nightly + @pytest.mark.precommit + @pytest.mark.parametrize('p', [-1, 1, float('inf'), float('-inf'), "fro"]) + @pytest.mark.parametrize('dim', [[0, 1], [-1, -2]]) + @pytest.mark.parametrize('keepdim', [True, False]) + @pytest.mark.parametrize("dtype", ["float32", "float64", None]) + @pytest.mark.parametrize("out", [True, False]) + @pytest.mark.parametrize("prim_dtype", [True, False]) + def test_linalg_matrix_norm(self, p, dim, keepdim, dtype, out, prim_dtype, ie_device, precision, ir_version): + self._test(*self.create_model(p, dim, keepdim, dtype, out, prim_dtype), + ie_device, precision, ir_version, + kwargs_to_prepare_input={"out": out or prim_dtype, "out_dtype": dtype if prim_dtype else None}) + + +class TestLinalgNorm(PytorchLayerTest): + + def _prepare_input(self, out=False, out_dtype=None): + if not out: + return (np.random.randn(3, 3).astype(np.float32),) + x = np.random.randn(3, 3).astype(np.float32) + y = np.random.randn(3, 3).astype( + out_dtype if out_dtype is not None else np.float32) + return (x, y) + + def create_model(self, p, dim, keepdim, dtype_str, out, out_as_dtype): + dtypes = { + "float32": torch.float32, + "float64": torch.float64 + } + dtype = dtypes.get(dtype_str) + + class aten_linalg_matrix_norm(torch.nn.Module): + + def __init__(self, p, dim, keepdim, dtype, out, out_as_dtype) -> None: + super().__init__() + self.ord = p + self.dim = dim + self.keepdim = keepdim + self.dtype = dtype + if self.dtype is not None: + self.forward = self.forward_dtype + if out: + self.forward = self.forward_out + if out_as_dtype: + self.forward = self.forward_prim_dtype + + def forward(self, x): + return torch.linalg.norm( + x, ord=self.ord, dim=self.dim, keepdim=self.keepdim + ) + + def forward_dtype(self, x): + return torch.linalg.norm( + x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=self.dtype + ) + + def forward_prim_dtype(self, x, y): + return torch.linalg.norm( + x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, dtype=y.dtype + ) + + def forward_out(self, x, y): + return y, torch.linalg.norm( + x, ord=self.ord, dim=self.dim, keepdim=self.keepdim, out=y + ) + + ref_net = None + + return aten_linalg_matrix_norm(p, dim, keepdim, dtype, out, out_as_dtype), ref_net, "aten::linalg_norm" + + @pytest.mark.nightly + @pytest.mark.precommit + @pytest.mark.parametrize('p,dim', [ + (-1, [0, 1]), (1, [-1, -2]), (float('inf'), [1, 0]), + (float('-inf'), [-2, -1]), (0, 1), (1, -1), + (None, None), (2.5, 0), (-1, 1), (2, 0), + (float('inf'), 1), (float('-inf'), 1), ("fro", (0, 1))]) + @pytest.mark.parametrize('keepdim', [True, False]) + @pytest.mark.parametrize("dtype", ["float32", "float64", None]) + @pytest.mark.parametrize("out", [True, False]) + @pytest.mark.parametrize("prim_dtype", [True, False]) + def test_linalg_norm(self, p, dim, keepdim, dtype, out, prim_dtype, ie_device, precision, ir_version): + self._test(*self.create_model(p, dim, keepdim, dtype, out, prim_dtype), + ie_device, precision, ir_version, + kwargs_to_prepare_input={"out": out or prim_dtype, "out_dtype": dtype if prim_dtype else None})