Skip to content

Commit 44967ef

Browse files
mvafinmvafin
authored
[ONNX] Support reshaping LSTM input (#34657)
### Details: - *Support reshaping LSTM input that is missing `num_directions` dimension* ### Tickets: - *CVS-182578* ### AI Assistance: - *AI assistance used: no / yes* - *If yes, summarize how AI was used and what human validation was performed (build/tests/manual checks).*
1 parent 2ed14e4 commit 44967ef

File tree

3 files changed

+264
-12
lines changed

3 files changed

+264
-12
lines changed

‎src/frontends/onnx/frontend/src/op/lstm.cpp‎

Lines changed: 60 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#include "openvino/op/multiply.hpp"
1515
#include "openvino/op/shape_of.hpp"
1616
#include "openvino/op/squeeze.hpp"
17+
#include "openvino/op/unsqueeze.hpp"
1718
#include "openvino/util/common_util.hpp"
1819
#include "utils/reshape.hpp"
1920
#include "utils/split.hpp"
@@ -38,10 +39,12 @@ enum class LSTMInput {
3839
LSTM_INPUT_P
3940
};
4041

41-
// Normalize tensor rank to target_rank by squeezing leading dimensions of size 1.
42-
// This handles models where upstream Unsqueeze operations add extra leading dimensions.
43-
// Static validation: if leading dims are statically known and != 1, emit clear error at conversion time.
44-
// If extra dimensions are != 1 at runtime, Squeeze will fail with a clear error.
42+
// Normalize tensor rank to target_rank:
43+
// - If rank > target: squeeze leading dimensions of size 1 (handles upstream Unsqueeze ops).
44+
// Static validation rejects leading dims that are statically != 1.
45+
// - If rank < target by exactly 1: unsqueeze a leading dimension of size 1
46+
// (handles models that omit the num_directions=1 dimension for unidirectional LSTMs).
47+
// - If rank differs by more than the above, throw an error for malformed input.
4548
ov::Output<ov::Node> normalize_tensor_rank(const ov::Output<ov::Node>& input,
4649
int64_t target_rank,
4750
const std::string& input_name) {
@@ -86,14 +89,23 @@ ov::Output<ov::Node> normalize_tensor_rank(const ov::Output<ov::Node>& input,
8689
return std::make_shared<v0::Squeeze>(input, axes_const);
8790
}
8891

89-
// input_rank < target_rank: reject non-conformant models with missing num_directions dimension
90-
OPENVINO_THROW("LSTM input '",
91-
input_name,
92-
"' has rank ",
93-
input_rank.get_length(),
94-
" but expected ",
95-
target_rank,
96-
". Missing num_directions dimension cannot be automatically inferred.");
92+
// input_rank < target_rank: only allow exactly 1 missing dimension (the num_directions dim).
93+
// For unidirectional LSTM (forward/reverse), num_directions=1, so some models omit it.
94+
// A larger rank deficiency indicates a genuinely malformed model.
95+
const auto dims_to_unsqueeze = target_rank - input_rank.get_length();
96+
if (dims_to_unsqueeze != 1) {
97+
OPENVINO_THROW("LSTM input '",
98+
input_name,
99+
"' has rank ",
100+
input_rank.get_length(),
101+
" but expected ",
102+
target_rank,
103+
". Rank difference is ",
104+
dims_to_unsqueeze,
105+
" but only 1 (missing num_directions) is supported.");
106+
}
107+
auto axes_const = v0::Constant::create(ov::element::i64, Shape{1}, std::vector<int64_t>{0});
108+
return std::make_shared<v0::Unsqueeze>(input, axes_const);
97109
}
98110

99111
struct LSTMNgInputMap {
@@ -116,6 +128,24 @@ struct LSTMNgInputMap {
116128

117129
m_input_map[LSTMInput::LSTM_INPUT_X] = input_x;
118130

131+
// Detect if num_directions dimension is missing from W.
132+
// Some models omit the leading num_directions dimension when it equals 1.
133+
// normalize_tensor_rank will unsqueeze it, but we need to squeeze the
134+
// corresponding dimension from outputs to match the original model's expectations.
135+
const auto& w_rank = ng_inputs.at(1).get_partial_shape().rank();
136+
if (w_rank.is_static() && w_rank.get_length() < 3) {
137+
// Unsqueezing adds num_directions=1, which is only valid for forward/reverse.
138+
// For bidirectional LSTMs, num_directions=2 and cannot be inferred.
139+
const std::string direction =
140+
ov::util::to_lower(node.get_attribute_value<std::string>("direction", "forward"));
141+
OPENVINO_ASSERT(direction != "bidirectional",
142+
"LSTM input 'W' has rank ",
143+
w_rank.get_length(),
144+
" but expected 3. Cannot add num_directions dimension for bidirectional LSTM "
145+
"because num_directions=2 cannot be inferred from the data.");
146+
m_num_directions_unsqueezed = true;
147+
}
148+
119149
// Weight tensor for the gates.
120150
// ONNX Shape: [num_directions, 4*hidden_size, input_size]
121151
auto input_w = normalize_tensor_rank(ng_inputs.at(1), 3, "W");
@@ -252,6 +282,9 @@ struct LSTMNgInputMap {
252282
return m_input_map.at(key);
253283
}
254284
std::map<LSTMInput, ov::Output<ov::Node>> m_input_map;
285+
// True when num_directions dimension was missing from inputs and was added via Unsqueeze.
286+
// In this case, outputs need the num_directions dimension squeezed to match the original model.
287+
bool m_num_directions_unsqueezed = false;
255288
};
256289

257290
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ATTRIBUTES PARSING ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -308,6 +341,21 @@ ov::OutputVector lstm(const ov::frontend::onnx::Node& node) {
308341
const auto Y_h = lstm_sequence->output(1);
309342
const auto Y_c = lstm_sequence->output(2);
310343

344+
if (input_map.m_num_directions_unsqueezed) {
345+
// The num_directions dimension was added to inputs via Unsqueeze.
346+
// Squeeze it from outputs so downstream consumers see the original ranks.
347+
// Y: OV [batch_size, num_directions(1), seq_length, hidden_size] -> ONNX [seq_length, batch_size, hidden_size]
348+
// Y_h: OV [batch_size, num_directions(1), hidden_size] -> ONNX [batch_size, hidden_size]
349+
// Y_c: OV [batch_size, num_directions(1), hidden_size] -> ONNX [batch_size, hidden_size]
350+
auto num_dir_axis = v0::Constant::create(ov::element::i64, Shape{1}, {1});
351+
auto Y_squeezed = std::make_shared<v0::Squeeze>(Y, num_dir_axis);
352+
auto Y_h_squeezed = std::make_shared<v0::Squeeze>(Y_h, num_dir_axis);
353+
auto Y_c_squeezed = std::make_shared<v0::Squeeze>(Y_c, num_dir_axis);
354+
355+
// Y: [batch_size, seq_length, hidden_size] -> [seq_length, batch_size, hidden_size]
356+
return {ov::op::util::reorder_axes(Y_squeezed, {1, 0, 2}), Y_h_squeezed, Y_c_squeezed};
357+
}
358+
311359
return {ov::op::util::reorder_axes(Y, {2, 1, 0, 3}),
312360
ov::op::util::reorder_axes(Y_h, {1, 0, 2}),
313361
ov::op::util::reorder_axes(Y_c, {1, 0, 2})};

‎src/frontends/onnx/tests/models/lstm_missing_num_directions.prototxt‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
ir_version: 7
2+
producer_name: "OpenVINO Test"
3+
graph {
4+
node {
5+
input: "X"
6+
input: "W"
7+
input: "R"
8+
input: "B"
9+
output: "Y"
10+
output: "Y_h"
11+
output: "Y_c"
12+
op_type: "LSTM"
13+
attribute {
14+
name: "direction"
15+
s: "forward"
16+
type: STRING
17+
}
18+
attribute {
19+
name: "hidden_size"
20+
i: 2
21+
type: INT
22+
}
23+
}
24+
name: "lstm_missing_num_directions"
25+
input {
26+
name: "X"
27+
type {
28+
tensor_type {
29+
elem_type: 1
30+
shape {
31+
dim { dim_value: 3 }
32+
dim { dim_value: 2 }
33+
dim { dim_value: 4 }
34+
}
35+
}
36+
}
37+
}
38+
input {
39+
name: "W"
40+
type {
41+
tensor_type {
42+
elem_type: 1
43+
shape {
44+
dim { dim_value: 8 }
45+
dim { dim_value: 4 }
46+
}
47+
}
48+
}
49+
}
50+
input {
51+
name: "R"
52+
type {
53+
tensor_type {
54+
elem_type: 1
55+
shape {
56+
dim { dim_value: 8 }
57+
dim { dim_value: 2 }
58+
}
59+
}
60+
}
61+
}
62+
input {
63+
name: "B"
64+
type {
65+
tensor_type {
66+
elem_type: 1
67+
shape {
68+
dim { dim_value: 16 }
69+
}
70+
}
71+
}
72+
}
73+
output {
74+
name: "Y"
75+
type {
76+
tensor_type {
77+
elem_type: 1
78+
}
79+
}
80+
}
81+
output {
82+
name: "Y_h"
83+
type {
84+
tensor_type {
85+
elem_type: 1
86+
}
87+
}
88+
}
89+
output {
90+
name: "Y_c"
91+
type {
92+
tensor_type {
93+
elem_type: 1
94+
}
95+
}
96+
}
97+
}
98+
opset_import {
99+
domain: ""
100+
version: 7
101+
}

‎src/frontends/onnx/tests/onnx_import_rnn.in.cpp‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -609,6 +609,109 @@ OPENVINO_TEST(${BACKEND_NAME}, onnx_model_lstm_rank4_with_unsqueeze) {
609609
test_case.run_with_tolerance_as_fp(1.0e-4f);
610610
}
611611

612+
// Test for LSTM with inputs missing the num_directions dimension.
613+
// Some models omit the leading num_directions=1 dimension from W, R, B.
614+
// The converter should unsqueeze it automatically and squeeze it from outputs.
615+
// W shape: [8, 4] instead of [1, 8, 4], R: [8, 2] instead of [1, 8, 2], B: [16] instead of [1, 16]
616+
// Expected output Y shape: [3, 2, 2] (no num_directions), not [3, 1, 2, 2]
617+
OPENVINO_TEST(${BACKEND_NAME}, onnx_model_lstm_missing_num_directions) {
618+
auto model = convert_model("lstm_missing_num_directions.onnx");
619+
auto test_case = ov::test::TestCase(model, s_device);
620+
621+
// X: [seq_length=3, batch_size=2, input_size=4]
622+
std::vector<float> X = {0.49671414494514465f, -0.13826429843902588f, 0.6476885676383972f, 1.5230298042297363f,
623+
-0.2341533750295639f, -0.23413695394992828f, 1.5792127847671509f, 0.7674347162246704f,
624+
-0.4694743752479553f, 0.5425600409507751f, -0.4634176790714264f, -0.4657297432422638f,
625+
0.241962268948555f, -1.9132802486419678f, -1.7249178886413574f, -0.5622875094413757f,
626+
-1.0128310918807983f, 0.31424733996391296f, -0.9080240726470947f, -1.4123036861419678f,
627+
1.4656487703323364f, -0.2257762998342514f, 0.06752820312976837f, -1.424748182296753f};
628+
629+
// W: [4*hidden_size=8, input_size=4] - missing num_directions dimension
630+
std::vector<float> W = {-0.5443827509880066f, 0.11092258989810944f, -1.1509935855865479f, 0.3756980299949646f,
631+
-0.6006386876106262f, -0.2916937470436096f, -0.6017066240310669f, 1.852278232574463f,
632+
-0.013497225008904934f, -1.057710886001587f, 0.8225449323654175f, -1.2208436727523804f,
633+
0.20886360108852386f, -1.959670066833496f, -1.32818603515625f, 0.19686123728752136f,
634+
0.7384665608406067f, 0.1713682860136032f, -0.1156482845544815f, -0.3011036813259125f,
635+
-1.4785219430923462f, -0.7198442220687866f, -0.46063876152038574f, 1.0571222305297852f,
636+
0.3436183035373688f, -1.7630401849746704f, 0.32408398389816284f, -0.38508227467536926f,
637+
-0.6769220232963562f, 0.6116762757301331f, 1.0309995412826538f, 0.9312801361083984f};
638+
639+
// R: [4*hidden_size=8, hidden_size=2] - missing num_directions dimension
640+
std::vector<float> R = {-0.8392175436019897f,
641+
-0.3092123866081238f,
642+
0.3312634229660034f,
643+
0.9755451083183289f,
644+
-0.4791742265224457f,
645+
-0.18565897643566132f,
646+
-1.106334924697876f,
647+
-1.1962065696716309f,
648+
0.8125258088111877f,
649+
1.3562400341033936f,
650+
-0.07201012223958969f,
651+
1.003532886505127f,
652+
0.3616360127925873f,
653+
-0.6451197266578674f,
654+
0.36139559745788574f,
655+
1.538036584854126f};
656+
657+
// B: [8*hidden_size=16] - missing num_directions dimension
658+
std::vector<float> B = {-0.03582603856921196f,
659+
1.5646436214447021f,
660+
-2.6197450160980225f,
661+
0.8219025135040283f,
662+
0.08704707026481628f,
663+
-0.2990073561668396f,
664+
0.0917607769370079f,
665+
-1.9875688552856445f,
666+
-0.21967189013957977f,
667+
0.3571125566959381f,
668+
1.4778940677642822f,
669+
-0.5182701945304871f,
670+
-0.8084936141967773f,
671+
-0.501757025718689f,
672+
0.9154021143913269f,
673+
0.3287511169910431f};
674+
675+
// Expected outputs - verified against reference model with proper [1, ...] shapes
676+
// Y: [seq_length=3, batch_size=2, hidden_size=2] - no num_directions dim
677+
std::vector<float> expected_Y = {0.022259337827563286f,
678+
0.003388261189684272f,
679+
0.05276688188314438f,
680+
0.11502056568861008f,
681+
0.00648046750575304f,
682+
-0.2669661343097687f,
683+
0.35137197375297546f,
684+
-0.5121785998344421f,
685+
0.13379308581352234f,
686+
-0.4288314878940582f,
687+
0.37510907649993896f,
688+
-0.0907968208193779f};
689+
690+
// Y_h: [batch_size=2, hidden_size=2] - no num_directions dim
691+
std::vector<float> expected_Y_h = {0.13379308581352234f,
692+
-0.4288314878940582f,
693+
0.37510907649993896f,
694+
-0.0907968208193779f};
695+
696+
// Y_c: [batch_size=2, hidden_size=2] - no num_directions dim
697+
std::vector<float> expected_Y_c = {0.35507577657699585f,
698+
-0.7553168535232544f,
699+
0.6096105575561523f,
700+
-0.12818995118141174f};
701+
702+
test_case.add_input<float>(Shape{3, 2, 4}, X);
703+
test_case.add_input<float>(Shape{8, 4}, W);
704+
test_case.add_input<float>(Shape{8, 2}, R);
705+
test_case.add_input<float>(Shape{16}, B);
706+
707+
// Output shapes lack the num_directions dimension
708+
test_case.add_expected_output<float>(Shape{3, 2, 2}, expected_Y);
709+
test_case.add_expected_output<float>(Shape{2, 2}, expected_Y_h);
710+
test_case.add_expected_output<float>(Shape{2, 2}, expected_Y_c);
711+
712+
test_case.run_with_tolerance_as_fp(1.0e-4f);
713+
}
714+
612715
// RNNLikeSequenceOp test fixture for test setup reuse
613716
class GRUSequenceOp : public testing::Test {
614717
public:

0 commit comments

Comments
 (0)