From 4be6025fae8c29c23f85846d700ff0d15381e11b Mon Sep 17 00:00:00 2001 From: Maxim Vafin Date: Mon, 10 Jun 2024 13:35:58 +0200 Subject: [PATCH] [PT FE] Fix segfault when prim::If can't be converted (#24914) ### Details: - *When scripting upsample creates `prim::If` depending from input shape, to determine which version of upsample to use: 1d, 2d or 3d. This is not supported by openvino, but this PR fixes segfault that happen in this case.* ### Tickets: - *#24270* --- .../pytorch/src/pt_framework_node.hpp | 5 ++- src/frontends/pytorch/src/utils.cpp | 8 ++-- .../pytorch_tests/test_upsample.py | 37 ++++++++++++++++++- 3 files changed, 44 insertions(+), 6 deletions(-) diff --git a/src/frontends/pytorch/src/pt_framework_node.hpp b/src/frontends/pytorch/src/pt_framework_node.hpp index c8346a6d48b..dddac5cde22 100644 --- a/src/frontends/pytorch/src/pt_framework_node.hpp +++ b/src/frontends/pytorch/src/pt_framework_node.hpp @@ -20,8 +20,9 @@ public: PtFrameworkNode(const std::shared_ptr& decoder, const OutputVector& inputs, size_t output_size, - bool is_reverseprop = false) - : ov::op::util::FrameworkNode(inputs, output_size, decoder->get_subgraph_size()), + bool is_reverseprop = false, + bool skip_subgraphs = false) + : ov::op::util::FrameworkNode(inputs, output_size, skip_subgraphs ? 0 : decoder->get_subgraph_size()), m_decoder(decoder) { ov::op::util::FrameworkNodeAttrs attrs; attrs.set_type_name("PTFrameworkNode"); diff --git a/src/frontends/pytorch/src/utils.cpp b/src/frontends/pytorch/src/utils.cpp index 0e6d0535763..c233c5ae297 100644 --- a/src/frontends/pytorch/src/utils.cpp +++ b/src/frontends/pytorch/src/utils.cpp @@ -213,8 +213,9 @@ namespace { std::shared_ptr create_fw_node_with_exception(const NodeContext& context, const ov::OutputVector& inputs, size_t num_outputs, - const std::string& exception_message) { - auto fw_node = std::make_shared(context.get_decoder(), inputs, num_outputs); + const std::string& exception_message, + bool skip_subgraphs = false) { + auto fw_node = std::make_shared(context.get_decoder(), inputs, num_outputs, false, skip_subgraphs); context.mark_node(fw_node); auto attrs = fw_node->get_attrs(); std::string message(exception_message); @@ -229,7 +230,8 @@ std::shared_ptr create_fw_node_with_exception(const NodeContext } // namespace OutputVector make_framework_node_ignore_bodies(const NodeContext& context, const std::string& exception) { - auto fw_node = create_fw_node_with_exception(context, context.inputs(), context.get_output_size() + 1, exception); + auto fw_node = + create_fw_node_with_exception(context, context.inputs(), context.get_output_size() + 1, exception, true); return fw_node->outputs(); } diff --git a/tests/layer_tests/pytorch_tests/test_upsample.py b/tests/layer_tests/pytorch_tests/test_upsample.py index 8811c30f965..34ffb9880c7 100644 --- a/tests/layer_tests/pytorch_tests/test_upsample.py +++ b/tests/layer_tests/pytorch_tests/test_upsample.py @@ -214,4 +214,39 @@ class TestUpsample2DListSizes(PytorchLayerTest): @pytest.mark.nightly @pytest.mark.precommit def test_upsample2d_list_sizes(self, mode, ie_device, precision, ir_version): - self._test(*self.create_model(mode), ie_device, precision, ir_version, trace_model=True) + self._test(*self.create_model(mode), ie_device, + precision, ir_version, trace_model=True) + + +class TestUpsampleScripted(PytorchLayerTest): + def _prepare_input(self): + import numpy as np + return (np.random.randn(1, 3, 200, 200).astype(np.float32),) + + def create_model(self): + import torch.nn as nn + + class TestModel(nn.Module): + def __init__(self, n_channels, n_classes): + super().__init__() + self.n_channels = n_channels + self.n_classes = n_classes + + self.cv1 = nn.Conv2d(n_channels, 16, kernel_size=3, padding=1) + self.mp1 = nn.MaxPool2d((2, 2), (2, 2)) + self.up = nn.Upsample(scale_factor=2.) + + def forward(self, x): + x1 = self.cv1(x) + x2 = self.mp1(x1) + x3 = self.up(x2) + return x3 + + return TestModel(1, 3), None, ["prim::If", "aten::upsample_nearest1d", "aten::upsample_nearest2d", "aten::upsample_nearest3d"] + + @pytest.mark.nightly + @pytest.mark.precommit + @pytest.mark.xfail(reason="Scripted upsample is not supported") + def test_upsample_scripted(self, ie_device, precision, ir_version): + self._test(*self.create_model(), ie_device, + precision, ir_version, trace_model=False)