From 2c88fbf798b43a16bceed64430f6d4d0871a2ac5 Mon Sep 17 00:00:00 2001 From: Ekaterina Aidova Date: Wed, 20 Sep 2023 11:26:18 +0400 Subject: [PATCH] [PT FE]: support mixed precision in aten::min/max (#19936) * [PT FE]: support mixed precision in aten::min/max * fix eltwise dtype alignment for float16 --- src/frontends/pytorch/src/op/min_max.cpp | 2 + src/frontends/pytorch/src/utils.cpp | 8 ++-- .../pytorch_tests/pytorch_layer_test_class.py | 7 ++-- tests/layer_tests/pytorch_tests/test_add.py | 6 +++ tests/layer_tests/pytorch_tests/test_div.py | 6 +++ .../layer_tests/pytorch_tests/test_min_max.py | 40 ++++++++++++++----- tests/layer_tests/pytorch_tests/test_mul.py | 6 +++ 7 files changed, 57 insertions(+), 18 deletions(-) diff --git a/src/frontends/pytorch/src/op/min_max.cpp b/src/frontends/pytorch/src/op/min_max.cpp index ad5c8224d3e..c6b63f11e25 100644 --- a/src/frontends/pytorch/src/op/min_max.cpp +++ b/src/frontends/pytorch/src/op/min_max.cpp @@ -33,6 +33,7 @@ OutputVector translate_max(const NodeContext& context) { // torch.max(input, other) if (context.input_is_none(2)) { auto y = context.get_input(1); + align_eltwise_input_types(context, x, y, true); return {context.mark_node(std::make_shared(x, y))}; } // torch.max(input, dim, keepdim), returns values and indicies @@ -62,6 +63,7 @@ OutputVector translate_min(const NodeContext& context) { // torch.min(input, other) if (context.input_is_none(2)) { auto y = context.get_input(1); + align_eltwise_input_types(context, x, y, true); return {context.mark_node(std::make_shared(x, y))}; } // torch.min(input, dim, keepdim), returns values and indicies diff --git a/src/frontends/pytorch/src/utils.cpp b/src/frontends/pytorch/src/utils.cpp index 386ac393d35..afd9b5ebf85 100644 --- a/src/frontends/pytorch/src/utils.cpp +++ b/src/frontends/pytorch/src/utils.cpp @@ -422,9 +422,9 @@ void align_eltwise_input_types(const NodeContext& context, Output& lhs, Ou // if div we need to also align float types to highest bitness regardless of scalar if (!align_scalars) lhs_dst_type = element::f32; - rhs_dst_type = element::f32; + rhs_dst_type = lhs_type; } else if (is_rhs_scalar && !lhs_type.is_real() && rhs_type.is_real()) { - lhs_dst_type = element::f32; + lhs_dst_type = rhs_type; // if div we need to also align float types to highest bitness regardless of scalar if (!align_scalars) rhs_dst_type = element::f32; @@ -437,9 +437,9 @@ void align_eltwise_input_types(const NodeContext& context, Output& lhs, Ou } if (!lhs_dst_type.is_real() && rhs_dst_type.is_real()) { - lhs_dst_type = element::f32; + lhs_dst_type = rhs_dst_type; } else if (lhs_dst_type.is_real() && !rhs_dst_type.is_real()) { - rhs_dst_type = element::f32; + rhs_dst_type = lhs_dst_type; } // Align bool to other type if (lhs_dst_type == element::boolean) { diff --git a/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py b/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py index 0f5638ea8c8..00cdc427056 100644 --- a/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py +++ b/tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py @@ -25,7 +25,8 @@ class PytorchLayerTest: "int64": Type.i64, "int16": Type.i16, "int8": Type.i8, - "uint8": Type.u8 + "uint8": Type.u8, + "float16": Type.f16 } @staticmethod @@ -120,8 +121,8 @@ class PytorchLayerTest: continue assert ov_type == fw_type, f"dtype validation failed: {ov_type} != {fw_type}" continue - assert torch.tensor(np.array( - ov_tensor)).dtype == fw_tensor.dtype, f"dtype validation failed: {torch.tensor(np.array(ov_tensor)).dtype} != {fw_tensor.dtype}" + ov_tensor_fw_format = torch.tensor(np.array(ov_tensor)) + assert ov_tensor_fw_format.dtype == fw_tensor.dtype, f"dtype validation failed: {ov_tensor_fw_format.dtype} != {fw_tensor.dtype}" # Compare Ie results with Framework results fw_eps = custom_eps if precision == 'FP32' else 5e-2 diff --git a/tests/layer_tests/pytorch_tests/test_add.py b/tests/layer_tests/pytorch_tests/test_add.py index 8c3026a9c2c..7cf243b0577 100644 --- a/tests/layer_tests/pytorch_tests/test_add.py +++ b/tests/layer_tests/pytorch_tests/test_add.py @@ -96,6 +96,12 @@ class TestAddTypes(PytorchLayerTest): [torch.float32, torch.int32], [torch.float32, torch.int64], [torch.float32, torch.float64], + [torch.float16, torch.uint8], + [torch.uint8, torch.float16], + [torch.float16, torch.int32], + [torch.int32, torch.float16], + [torch.float16, torch.int64], + [torch.int64, torch.float16] ]) @pytest.mark.parametrize(("lhs_shape", "rhs_shape"), [([2, 3], [2, 3]), ([2, 3], []), diff --git a/tests/layer_tests/pytorch_tests/test_div.py b/tests/layer_tests/pytorch_tests/test_div.py index 48eaba1a89b..d6e696b6288 100644 --- a/tests/layer_tests/pytorch_tests/test_div.py +++ b/tests/layer_tests/pytorch_tests/test_div.py @@ -98,6 +98,12 @@ class TestDivTypes(PytorchLayerTest): [torch.float32, torch.int32], [torch.float32, torch.int64], [torch.float32, torch.float64], + [torch.float16, torch.uint8], + [torch.uint8, torch.float16], + [torch.float16, torch.int32], + [torch.int32, torch.float16], + [torch.float16, torch.int64], + [torch.int64, torch.float16] ]) @pytest.mark.parametrize(("lhs_shape", "rhs_shape"), [([2, 3], [2, 3]), ([2, 3], []), diff --git a/tests/layer_tests/pytorch_tests/test_min_max.py b/tests/layer_tests/pytorch_tests/test_min_max.py index 211cf76d06a..1d03c837380 100644 --- a/tests/layer_tests/pytorch_tests/test_min_max.py +++ b/tests/layer_tests/pytorch_tests/test_min_max.py @@ -7,19 +7,27 @@ from pytorch_layer_test_class import PytorchLayerTest class TestMinMax(PytorchLayerTest): - def _prepare_input(self, second_input=False): + def _prepare_input(self, input_dtype="float32", second_input=False, second_input_dtype="float32"): import numpy as np if not second_input: - return (np.random.randn(1, 3, 10, 10).astype(np.float32),) - return (np.random.randn(1, 3, 10, 10).astype(np.float32), np.random.randn(1, 3, 10, 10).astype(np.float32)) + return (np.random.randn(1, 3, 10, 10).astype(input_dtype),) + return (np.random.randn(1, 3, 10, 10).astype(input_dtype), np.random.randn(1, 3, 10, 10).astype(second_input_dtype)) - def create_model(self, op_type, axes, keep_dims, single_input=True): + def create_model(self, op_type, axes, keep_dims, single_input=True, dtypes=("float32", "float32")): import torch op_types = { 'max': torch.max, 'min': torch.min } + dtypes_map = { + "float32": torch.float32, + "float64": torch.float64, + "int32": torch.int32, + "int64": torch.int64, + "uint8": torch.uint8 + } + op = op_types[op_type] class aten_min_max(torch.nn.Module): @@ -41,17 +49,23 @@ class TestMinMax(PytorchLayerTest): return self.op(x, self.axes, self.keep_dims) class aten_min_max_2args(torch.nn.Module): - def __init__(self, op): + def __init__(self, op, l_dtype, r_dtype): super(aten_min_max_2args, self).__init__() self.op = op + self.l_dtype = l_dtype + self.r_dtype = r_dtype def forward(self, x, y): - return self.op(x, y) + return self.op(x.to(self.l_dtype), y.to(self.r_dtype)) ref_net = None if axes is None and keep_dims is None: - model_cls = aten_min_max( - op) if single_input else aten_min_max_2args(op) + if single_input: + model_cls = aten_min_max(op) + else: + l_dtype = dtypes_map[dtypes[0]] + r_dtype = dtypes_map[dtypes[1]] + model_cls = aten_min_max_2args(op, l_dtype, r_dtype) else: model_cls = aten_min_max_3args(op, axes, keep_dims) @@ -66,11 +80,15 @@ class TestMinMax(PytorchLayerTest): single_input=True), ie_device, precision, ir_version) @pytest.mark.parametrize("op_type", ['min', 'max']) + @pytest.mark.parametrize("second_input_dtype", ["float32", "int32", "float64", "int64", "uint8"]) + @pytest.mark.parametrize("first_input_dtype", ["float32", "int32", "float64", "int64", "uint8"]) @pytest.mark.nightly @pytest.mark.precommit - def test_min_max(self, op_type, ie_device, precision, ir_version): - self._test(*self.create_model(op_type, None, None, single_input=False), - ie_device, precision, ir_version, kwargs_to_prepare_input={"second_input": True}) + def test_min_max(self, op_type, first_input_dtype, second_input_dtype, ie_device, precision, ir_version): + self._test(*self.create_model(op_type, None, None, single_input=False, dtypes=(first_input_dtype, second_input_dtype)), + ie_device, precision, ir_version, kwargs_to_prepare_input= + {"second_input": True, "input_dtype": first_input_dtype, "second_input_dtype": second_input_dtype} + ) class TestPrimMax(PytorchLayerTest): diff --git a/tests/layer_tests/pytorch_tests/test_mul.py b/tests/layer_tests/pytorch_tests/test_mul.py index 8e958f09569..f22fd7b03bc 100644 --- a/tests/layer_tests/pytorch_tests/test_mul.py +++ b/tests/layer_tests/pytorch_tests/test_mul.py @@ -87,6 +87,12 @@ class TestMulTypes(PytorchLayerTest): [torch.float32, torch.int32], [torch.float32, torch.int64], [torch.float32, torch.float64], + [torch.float16, torch.uint8], + [torch.uint8, torch.float16], + [torch.float16, torch.int32], + [torch.int32, torch.float16], + [torch.float16, torch.int64], + [torch.int64, torch.float16] ]) @pytest.mark.parametrize(("lhs_shape", "rhs_shape"), [([2, 3], [2, 3]), ([2, 3], []),