Error analysing LSTM model using X-Cube-AI: ""NOT IMPLEMENTED: Order of dimensions of input cannot be interpreted""

[Lemanoise]

As stated, I have problem analyzing the ONNX model ported from a pytorch model. The code of the model is simple as follows:

class LSTM(nn.Module):

    def __init__(self, input_size, hidden_size, num_layers, num_output=1):
        super(LSTM, self).__init__()
        
        self.num_layers = num_layers
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_output = num_output
        
        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size,
                            num_layers=num_layers, batch_first=True)
        
        self.fc = nn.Linear(hidden_size, num_output)
        
    def forward(self, x, device='cuda'):
        
        ula, (h_out, _) = self.lstm(x)
        
        out = self.fc(h_out[-1])
        
        return out

 Netron vis:

I will really appreciate any advice or workarounds. Thank you!

[Lemanoise]

Thank you very much!

[Julian E.]

Hello @Lemanoise,

I should know more about how to use LSTM properly with ST Edge AI Core on the start of next year.

But for the moment I managed to use more than one hidden size with pytorch by not using the batch_first option:

import torch
import torch.nn as nn

class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_output=1):
        super(LSTM, self).__init__()
        self.num_layers = num_layers
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_output = num_output

        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size,
                            num_layers=num_layers, batch_first=False)

        self.fc = nn.Linear(hidden_size, num_output)

    def forward(self, x, device='cuda'):

        ula, (h_out, _) = self.lstm(x)  
        out = self.fc(h_out[-1])  
        return out

# If not 1 : error
# other we can put anything
input_size = 10
hidden_size = 10
num_layers = 1
num_output = 10

model = LSTM(input_size, hidden_size, num_layers, num_output)

model.eval()

dummy_input = torch.randn(1, 128, input_size)  # Batch size 1, Sequence length 128, Input size 10

onnx_file_path = "simple_lstml.onnx"
torch.onnx.export(model, dummy_input, onnx_file_path, 
                  input_names=['input'], output_names=['output'])

It may help you

Have a good day,

Julian