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