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}")
Leave a Comment