diff --git a/src/frontends/pytorch/src/op/pad.cpp b/src/frontends/pytorch/src/op/pad.cpp index 4f6e1865995..02051a56e2a 100644 --- a/src/frontends/pytorch/src/op/pad.cpp +++ b/src/frontends/pytorch/src/op/pad.cpp @@ -126,6 +126,14 @@ OutputVector translate_constant_pad_nd_fx(const NodeContext& context) { return translate_pad_common(context, data, paddings, pad_value); } +OutputVector translate_reflection_pad_nd_fx(const NodeContext& context) { + num_inputs_check(context, 2, 2); + auto data = context.get_input(0); + auto paddings = context.const_input>(1); + Output pad_value = context.mark_node(v0::Constant::create(element::f32, Shape{}, {0})); + return translate_pad_common(context, data, paddings, pad_value, "reflect"); +} + } // 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 d853e7a3f67..5bd891fa37a 100644 --- a/src/frontends/pytorch/src/op_table.cpp +++ b/src/frontends/pytorch/src/op_table.cpp @@ -267,6 +267,7 @@ OP_CONVERTER(translate_max_dim_fx); OP_CONVERTER(translate_max_poolnd_fx); OP_CONVERTER(translate_mean_fx); OP_CONVERTER(translate_min_dim_fx); +OP_CONVERTER(translate_reflection_pad_nd_fx); OP_CONVERTER(translate_rsub_fx); OP_CONVERTER(translate_scalar_tensor_fx); OP_CONVERTER(translate_scaled_dot_product_attention_fx); @@ -845,6 +846,9 @@ const std::map get_supported_ops_fx() { {"aten.pixel_shuffle.default", op::translate_pixel_shuffle}, {"aten.pixel_unshuffle.default", op::translate_pixel_unshuffle}, {"aten.reciprocal.default", op::translate_reciprocal}, + {"aten.reflection_pad1d.default", op::translate_reflection_pad_nd_fx}, + {"aten.reflection_pad2d.default", op::translate_reflection_pad_nd_fx}, + {"aten.reflection_pad3d.default", op::translate_reflection_pad_nd_fx}, {"aten.relu.default", op::translate_1to1_match_1_inputs}, {"aten.relu_.default", op::inplace_op>}, {"aten.repeat.default", op::translate_1to1_match_2_inputs}, diff --git a/tests/layer_tests/pytorch_tests/test_pad.py b/tests/layer_tests/pytorch_tests/test_pad.py index d033e043efa..6f73061517c 100644 --- a/tests/layer_tests/pytorch_tests/test_pad.py +++ b/tests/layer_tests/pytorch_tests/test_pad.py @@ -57,7 +57,7 @@ class TestPad(PytorchLayerTest): @pytest.mark.precommit def test_pad4d(self, pads, mode, value, dtype, ie_device, precision, ir_version): self._test(*self.create_model(pads, mode, value), ie_device, precision, ir_version, - kwargs_to_prepare_input={'ndim': 4, "dtype": dtype}) + kwargs_to_prepare_input={"ndim": 4, "dtype": dtype}) @pytest.mark.parametrize("pads,mode,value,dtype", [ ((1, 2, 3, 4, 5, 6), "reflect", None, "float32"), @@ -88,7 +88,7 @@ class TestPad(PytorchLayerTest): @pytest.mark.nightly def test_pad5d(self, pads, mode, value, dtype, ie_device, precision, ir_version): self._test(*self.create_model(pads, mode, value), ie_device, precision, ir_version, - kwargs_to_prepare_input={'ndim': 5, "dtype": dtype}, trace_model=True) + kwargs_to_prepare_input={"ndim": 5, "dtype": dtype}, trace_model=True) @pytest.mark.parametrize("pads,mode,value,dtype", [ ((1, 2), "reflect", None, 'float32'), @@ -107,7 +107,7 @@ class TestPad(PytorchLayerTest): @pytest.mark.nightly def test_pad2d(self, pads, mode, value, dtype, ie_device, precision, ir_version): self._test(*self.create_model(pads, mode, value), ie_device, precision, ir_version, - kwargs_to_prepare_input={'ndim': 2, "dtype": dtype}, trace_model=True) + kwargs_to_prepare_input={"ndim": 2, "dtype": dtype}, trace_model=True) class TestPadListPaddingings(PytorchLayerTest): @@ -126,7 +126,7 @@ class TestPadListPaddingings(PytorchLayerTest): self.mode = mode self.value = value - def forward(self, x, pad_w:int, pad_h:int): + def forward(self, x, pad_w: int, pad_h: int): return F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2], value=self.value) ref_net = None @@ -154,7 +154,7 @@ class TestPadListPaddingings(PytorchLayerTest): @pytest.mark.precommit def test_pad4d(self, pad_w, pad_h, mode, value, dtype, ie_device, precision, ir_version): self._test(*self.create_model(mode, value), ie_device, precision, ir_version, - kwargs_to_prepare_input={'ndim': 4, "pad_w": pad_w, "pad_h": pad_h, "dtype": dtype}) + kwargs_to_prepare_input={"ndim": 4, "pad_w": pad_w, "pad_h": pad_h, "dtype": dtype}) @pytest.mark.parametrize("pad_w,pad_h,mode,value", [ (2, 0, "reflect", None), @@ -176,7 +176,7 @@ class TestPadListPaddingings(PytorchLayerTest): @pytest.mark.nightly def test_pad5d(self, pad_w, pad_h, mode, value, ie_device, precision, ir_version): self._test(*self.create_model(mode, value), ie_device, precision, ir_version, - kwargs_to_prepare_input={'ndim': 5, "pad_w": pad_w, "pad_h": pad_h}) + kwargs_to_prepare_input={"ndim": 5, "pad_w": pad_w, "pad_h": pad_h}) @pytest.mark.parametrize("pad_w,pad_h,mode,value,dtype", [ (2, 0, "reflect", None, "float32"), @@ -199,4 +199,48 @@ class TestPadListPaddingings(PytorchLayerTest): @pytest.mark.precommit def test_pad2d(self, pad_w, pad_h, mode, value, dtype, ie_device, precision, ir_version): self._test(*self.create_model(mode, value), ie_device, precision, ir_version, - kwargs_to_prepare_input={'ndim': 2, "pad_w": pad_w, "pad_h": pad_h, "dtype": dtype}) \ No newline at end of file + kwargs_to_prepare_input={"ndim": 2, "pad_w": pad_w, "pad_h": pad_h, "dtype": dtype}) + + +class TestReflectionPad(PytorchLayerTest): + def _prepare_input(self, ndim=4, dtype="float32"): + import numpy as np + input_5d_shape = [1, 3, 14, 14, 18] + return (np.random.randn(*input_5d_shape[:ndim]).astype(dtype),) + + def create_model(self, pads): + import torch + import torch.nn.functional as F + + class aten_pad(torch.nn.Module): + def __init__(self, pads): + super().__init__() + ndim = len(pads) / 2 + if ndim == 1: + self.pad = torch.nn.ReflectionPad1d(pads) + elif ndim == 2: + self.pad = torch.nn.ReflectionPad1d(pads) + elif ndim == 3: + self.pad = torch.nn.ReflectionPad1d(pads) + else: + raise Exception("Unsupported pads") + + def forward(self, x): + return self.pad(x) + + # it will be a reflection_pad in export, but not in TS + return aten_pad(pads), None, "aten::pad" + + @pytest.mark.parametrize("dtype", ["float32", "float64", "int32"]) + @pytest.mark.parametrize("pads", [ + (1, 2), + (1, 2, 3, 4), + (1, 2, 3, 4, 3, 2), + ]) + @pytest.mark.nightly + @pytest.mark.precommit_torch_export + def test_reflection_padnd(self, pads, dtype, ie_device, precision, ir_version): + ndim = len(pads) // 2 + 2 + print(ndim) + self._test(*self.create_model(pads), ie_device, precision, ir_version, + kwargs_to_prepare_input={"ndim": ndim, "dtype": dtype}) diff --git a/tests/model_hub_tests/pytorch/torch_utils.py b/tests/model_hub_tests/pytorch/torch_utils.py index 09a89104699..4a1ed3fd5f8 100644 --- a/tests/model_hub_tests/pytorch/torch_utils.py +++ b/tests/model_hub_tests/pytorch/torch_utils.py @@ -68,12 +68,15 @@ class TestTorchConvertModel(TestConvertModel): from torch.export import export from packaging import version from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder - import inspect - from openvino.frontend.pytorch.utils import prepare_example_inputs_and_model input_shapes = [] input_types = [] model_obj.eval() + # need to infer before export to initialize everything, otherwise it will be initialized with FakeTensors + if isinstance(self.example, dict): + pt_res = model_obj(**self.example) + else: + pt_res = model_obj(*self.example) if isinstance(self.example, dict): graph = export(model_obj, tuple(), self.example) for input_data in self.example.values(): @@ -105,10 +108,6 @@ class TestTorchConvertModel(TestConvertModel): if isinstance(self.example, dict): decoder._input_signature = list(self.example.keys()) ov_model = convert_model(decoder, example_input=self.example) - if isinstance(self.example, dict): - pt_res = model_obj(**self.example) - else: - pt_res = model_obj(*self.example) if isinstance(pt_res, dict): for i, k in enumerate(pt_res.keys()): ov_model.outputs[i].get_tensor().set_names({k})