Fix LSTM conversion for models with rank > 3 inputs from Unsqueeze operations #33023

evkotov · 2025-11-25T12:57:22Z

Details:

Subgraph from silero_vad.onnx

Problem

ONNX models exported from PyTorch frequently contain Unsqueeze operations before LSTM nodes. These operations add extra dimensions to tensors, resulting in rank-4 or rank-5 inputs to LSTM nodes. However, the ONNX LSTM specification strictly requires rank-3 inputs with shape [seq_length, batch_size, input_size]. Why this happens:
PyTorch models use various tensor shapes during training
During ONNX export, shape mismatches are "fixed" by inserting Unsqueeze nodes
These Unsqueeze operations add dimensions with size 1 to match expected shapes
The resulting LSTM inputs have rank > 3, violating ONNX LSTM specification

Real-world impact:

Models like silero_vad.onnx contain 4 LSTM nodes, all with Unsqueeze operations before them
Without this fix LSTM models fail to convert to OpenVINO IR

Solution
This fix adds automatic rank reduction in the ONNX Frontend LSTM converter (src/frontends/onnx/frontend/src/op/lstm.cpp). The implementation uses a two-strategy approach:

Squeeze Strategy (optimal path):
Used when all extra leading dimensions equal 1
Example: [1, 1, seq, batch, input] → [seq, batch, input]
Zero-cost operation that only changes metadata, no data movement
Applies to most real-world models (including silero_vad.onnx)
Reshape Strategy (fallback path):
Used when extra dimensions are > 1 or have dynamic shapes
Example: [2, 3, seq, batch, input] → [6, batch, input] (flattens leading dimensions)
Handles edge cases and dynamic shapes
Uses dynamic shape calculation at runtime

Implementation details:

New function reduce_tensor_rank() analyzes input tensor rank and shape
Automatically selects optimal strategy based on dimension values
Applied to all LSTM inputs: X (data), initial_h (hidden state), initial_c (cell state)
Transparent to users - no model modifications required
Code structure:

// Analyze input shape
if (input_rank <= target_rank) {
    return input;  // No reduction needed
}

// Check if all extra dimensions equal 1
if (all_extra_dims_are_one) {
    // Use Squeeze - optimal path
    return Squeeze(input, axes);
} else {
    // Use Reshape - fallback path
    return Reshape(input, new_shape);
}

Performance:
Squeeze path has zero runtime overhead (metadata-only operation)
Reshape path adds minimal overhead only for edge cases
No impact on models that already have rank-3 inputs

Tickets:

162986

Problem: Models like silero_vad contain LSTM layers with high-rank input tensors (rank > 3), but OpenVINO's LSTM expects exactly rank 3 [batch, sequence, features]. This causes conversion failures for models with shapes like [1, 1, ?, ?, ?]. Solution: Add reduce_tensor_rank() helper function that processes LSTM inputs (X, initial_h, initial_c) before axis reordering. The function: - Squeezes leading dimensions equal to 1 when possible - Uses Reshape to collapse leading dimensions if they aren't all 1 - Reduces rank to target rank 3 before reordering axes Test: Added onnx_model_lstm_high_rank_input test with rank-5 input [1,1,3,2,4] that gets reduced to [3,2,4]. Reference outputs generated using ONNX Runtime with equivalent rank-3 input.

This test does not reproduce silero_vad.onnx structure. Keep only tests that match actual silero_vad patterns: - lstm_rank5_squeeze: multi-axis squeeze (with GPU skip) - lstm_rank4_with_unsqueeze: exact silero_vad structure