[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*
This commit is contained in:
Maxim Vafin 2024-06-10 13:35:58 +02:00 committed by GitHub
parent 73a2d13721
commit 4be6025fae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 44 additions and 6 deletions

View File

@ -20,8 +20,9 @@ public:
PtFrameworkNode(const std::shared_ptr<TorchDecoder>& 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");

View File

@ -213,8 +213,9 @@ namespace {
std::shared_ptr<PtFrameworkNode> 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<PtFrameworkNode>(context.get_decoder(), inputs, num_outputs);
const std::string& exception_message,
bool skip_subgraphs = false) {
auto fw_node = std::make_shared<PtFrameworkNode>(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<PtFrameworkNode> 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();
}

View File

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