Untitled

import torch
import torch.nn as nn
import torch.onnx

# Define a single-layer LSTM model
class SingleLayerLSTM(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SingleLayerLSTM, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)  # Single-layer LSTM
        self.fc = nn.Linear(hidden_size, output_size)  # Fully connected layer for output

    def forward(self, x):
        # Pass through LSTM layer
        lstm_out, (hn, cn) = self.lstm(x)
        # Use the last hidden state directly (avoiding slicing)
        out = self.fc(hn[-1])  # hn[-1] corresponds to the last layer's hidden state
        return out

# Initialize the model
input_size = 10   # Number of input features
hidden_size = 20  # Number of hidden units in the LSTM
output_size = 5   # Number of output features
model = SingleLayerLSTM(input_size, hidden_size, output_size)
model.eval()  # Set the model to evaluation mode

# Create dummy input (batch size 1, sequence length 15, input size 10)
dummy_input = torch.randn(1, 15, input_size)

# Path to save the ONNX model
onnx_path = "single_layer_lstm_fixed.onnx"

# Export the model to ONNX
torch.onnx.export(
    model,                   # Model to export
    dummy_input,             # Example input
    onnx_path,               # Path to save the ONNX file
    export_params=True,      # Store trained parameter weights in the model
    opset_version=11,        # ONNX opset version
    input_names=["input"],   # Input tensor name
    output_names=["output"], # Output tensor name
    dynamic_axes={"input": {0: "batch_size", 1: "sequence_length"},  # Support dynamic input sizes
                  "output": {0: "batch_size"}}
)

print(f"Single-layer LSTM model (fixed) has been converted to ONNX and saved to {onnx_path}")