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

[PT FE] Support aten.reflection_pad{1,2,3}d.default (#23379) 

### Details:
- *Support `aten.reflection_pad1d.default`,
`aten.reflection_pad2d.default` and `aten.reflection_pad3d.default`*

### Tickets:
 - *ticket-id*
This commit is contained in:
Maxim Vafin 2024-03-13 23:39:37 +01:00 committed by GitHub
parent a139ccbba8
commit ad8d6ef35d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 68 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -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<std::vector<int64_t>>(1);
Output<Node> 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

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

@ -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<std::string, CreatorFunction> 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<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>},

58
tests/layer_tests/pytorch_tests/test_pad.py
View File

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

11
tests/model_hub_tests/pytorch/torch_utils.py
View File

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