Error analysing LSTM model using X-Cube-AI:INTERNAL ERROR : Unkonwn dimensions : CH

[Wangxingkun]

When I analyze the ONNX model in CubeAl,it says INTERNAL ERROR : Unkonwn dimensions : CH.The code of the model is simple as follows，input_size = 1 hidden_size = 1 output_size = 1 num_layers = 1
I don't know why and I will really appreciate any advice or workarounds. Thank you：

class LSTMNet(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=1)
        super(LSTMNet, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=False)  # batch_first=False
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        h0 = torch.zeros(num_layers, x.size(1), hidden_size).to(x.device) 
        c0 = torch.zeros(num_layers, x.size(1), hidden_size).to(x.device)
        out, _ = self.lstm(x, (h0, c0))
        out = self.fc(out[-1, :, :]) 
        return out

[Wangxingkun]

And this is my complete training code, if you know the reason why please reply me, thank you !!!

class SequenceDataset(Dataset):
    def __init__(self, sequences, labels):
        self.sequences = sequences
        self.labels = labels

    def __len__(self:(
        return len(self.labels)

    def __getitem__(self, idx):
        return torch.Tensor(self.sequences[idx]), torch.Tensor(self.labels[idx])

class LSTMNet(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=1:(
        super(LSTMNet, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=False)  # batch_first=False
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        h0 = torch.zeros(num_layers, x.size(1), hidden_size).to(x.device) 
        c0 = torch.zeros(num_layers, x.size(1), hidden_size).to(x.device)
        out, _ = self.lstm(x, (h0, c0))
        out = self.fc(out[-1, :, :]) 
        return out

input_size = 1 
hidden_size = 1 
output_size = 1
num_layers = 1
num_epochs = 120
batch_size = 32
learning_rate = 0.001

num_samples = 1000
sequence_length = 15
sequences = torch.rand(num_samples, sequence_length, input_size)
labels = torch.randint(0, 2, (num_samples, 1))

dataset = SequenceDataset(sequences.numpy(), labels.numpy())
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

model = LSTMNet(input_size, hidden_size, output_size, num_layers)
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

model.train()
for epoch in range(num_epochs):
    for seqs, lbls in train_loader:
        seqs = seqs.permute(1, 0, 2)  
        outputs = model(seqs)
        loss = criterion(outputs, lbls)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

example_input = torch.rand(sequence_length, 1, input_size) 
torch.onnx.export(model, example_input, 'lstm_model.onnx',
                  input_names=['input'],
                  output_names=['output'],
                  dynamic_axes={'input': {1: 'batch_size', 0: 'sequence_length'},
                                'output': {0: 'batch_size'}})
print('Model exported as lstm_model.onnx')

[Julian E.]

Hello @Wangxingkun,

I am not an expert regarding pytorch so I will need my colleagues for a more complete answer. But they are in vacation. In the meantime, you can take a look at this code that works and should help you

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'])

Have a good day,

Julian

[Wangxingkun]

Thank you very much Julian, I have tried your code but it still says 'INTERNAL ERROR : Unkonwn dimensions : CH'.It seems that as long as there is an LSTM layer in the model, there will be the following error. If it is a fully connected layer, there will be no such error.But in my project, the LSTM layer is necessary, and I don't know how to fix it.It makes me a little frustrated.

[Julian E. · Accepted solution]

Hello @Wangxingkun ,

What version of ST Edge AI are you using?

I tried it and I do not get errors...

This is the version I have (the last I believe)

Have a good day,

Julian

[Wangxingkun]

Hello Julian,

Thank you very much!! I tried your zip again and when I converted the version to 10.0.0, this error no longer occurred.From what I see now, this is just a minor version issue, but you still patiently helped me clarify my confusion and analyze possible problems. I am truly grateful to you.

Have a good day :) ,

Wang

[wypeng]

Dear Julian,

Thank you for your assistance and guidance. I would like to seek further advice regarding a challenge I encountered. When executing the provided model on my local environment, the predictions align well with expectations. However, when deploying the same model to an STM32 embedded board, the results differ.
Below is a simplified version of the Pytorch code and model for reference:

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

import numpy as np
from sklearn.preprocessing import MinMaxScaler

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

input_size = 1
hidden_size = 64
num_layers = 1
num_output = 1

model = LSTM(input_size, hidden_size, num_layers, num_output)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

seq_length = 50  
num_samples = 1000  
X, y = generate_sine_wave(seq_length, num_samples)
scaler = MinMaxScaler(feature_range=(0, 1))
X_scaled = scaler.fit_transform(X)
X_scaled = X_scaled.reshape((X_scaled.shape[0], X_scaled.shape[1], 1))
y = y.reshape(-1, 1)

# train
X_train = torch.tensor(X_scaled, dtype=torch.float32)  
y_train = torch.tensor(y, dtype=torch.float32)
train_loader = DataLoader(list(zip(X_train, y_train)), batch_size=1, shuffle=True)

epochs = 50
for epoch in range(epochs):
    model.train()
    for i, (X_batch, y_batch) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(X_batch)
        loss = criterion(output, y_batch)
        loss.backward()
        optimizer.step()
    print(f'Epoch: {epoch+1:2d}, loss: {loss.item()}')

X_batch = torch.randn(1, seq_length, input_size)
onnx_file_path = "torch_lstm.onnx"
torch.onnx.export(model, X_batch, onnx_file_path,
                  input_names=['input'], output_names=['output'])

Currently, I have confirmed that models built with Keras (without LSTM layers) produce consistent predictions between the local environment and the STM32. Below is a simplified version of the Keras code and model for reference:

from keras.models import Model
from keras.layers import Input, LSTM, Dense

input_layer = Input(shape=(seq_length, 1), name="input_layer")
lstm_1 = LSTM(32, activation='relu', return_sequences=True, name="lstm_1")(input_layer)
lstm_2 = LSTM(32, activation='relu', name="lstm_2")(lstm_1)
output_layer = Dense(units=1, activation='linear', name="output_layer")(lstm_2)
model = Model(inputs=input_layer, outputs=output_layer)

model.summary()
model.compile(optimizer='adam', loss='mse') #mean_squared_error
model.fit(X_scaled, y, epochs=30, batch_size=1, validation_split=0.2, 
          shuffle=True, verbose=1)

model.save('./keras_lstm.h5')

Below is a simplified test set for inference:

static float X_test_2d[50][1] = {

      {0.5908357501029968f},
      {0.6132590174674988f},
      {0.6401912569999695f},
      {0.6743371486663818f},
      {0.7207258343696594f},
      {0.7869139313697815f},
      {0.8616802096366882f},
      {0.9366850852966309f},
      {0.9935588240623474f},
      {0.9785177111625671f},
      {0.9057679772377014f},
      {0.82420414686203f},
      {0.7471164464950562f},
      {0.6801365613937378f},
      {0.6331720352172852f},
      {0.598699688911438f},
      {0.5715680122375488f},
      {0.549017608165741f},
      {0.5294104218482971f},
      {0.5116827487945557f},
      {0.4950788617134094f},
      {0.4790057837963104f},
      {0.4629416763782501f},
      {0.44636598229408264f},
      {0.42868947982788086f},
      {0.4091642498970032f},
      {0.3867409825325012f},
      {0.3598087728023529f},
      {0.32566285133361816f},
      {0.2792741358280182f},
      {0.2130860686302185f},
      {0.13831977546215057f},
      {0.06331492215394974f},
      {0.006441186182200909f},
      {0.021482301875948906f},
      {0.09423204511404037f},
      {0.17579583823680878f},
      {0.25288352370262146f},
      {0.3198634684085846f},
      {0.36682793498039246f},
      {0.401300311088562f},
      {0.42843198776245117f},
      {0.45098239183425903f},
      {0.4705895781517029f},
      {0.48831725120544434f},
      {0.5049211382865906f},
      {0.520994246006012f},
      {0.5370582938194275f},
      {0.5536340475082397f},
      {0.5713105201721191f}

};

static float y_test[1][1] = {
    {0.49757889651355536f}
};

When using PyTorch to build models and deploying them to an STM32 board, I noticed that the parameter size differ between the local environment and the STM32 after deployment, even with the same operations. However, when using Keras, the parameter size remain consistent between the two environments.

Could you kindly provide guidance on what I should pay attention to in order to resolve this issue?

Thank you for your time and support.

Best regards,

Wendy