### LSTM Implementation in Python (PyTorch) ```python import torch import torch.nn as nn class BasicLSTM(nn.Module): def __init__(self, input_size, hidden_size, output_size, num_layers=1): """ Args: input_size (int): The number of expected features in the input `x`. hidden_size (int): The number of features in the hidden state `h`. output_size (int): The number of output features. num_layers (int): Number of recurrent layers. """ super(BasicLSTM, self).__init__() self.hidden_size = hidden_size self.num_layers = num_layers # LSTM layer: input_size, hidden_size, num_layers, batch_first # batch_first=True means input/output tensors are (batch, seq, feature) self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True) # Fully connected layer to map LSTM output to desired output_size self.fc = nn.Linear(hidden_size, output_size) def forward(self, x): """ Args: x (torch.Tensor): Input tensor of shape (batch_size, sequence_length, input_size) Returns: torch.Tensor: Output tensor of shape (batch_size, output_size) """ # Initialize hidden state and cell state with zeros # h0 shape: (num_layers, batch_size, hidden_size) # c0 shape: (num_layers, batch_size, hidden_size) h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device) c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device) # Pass input through LSTM layer # out: (batch_size, sequence_length, hidden_size) # hn: (num_layers, batch_size, hidden_size) - final hidden state # cn: (num_layers, batch_size, hidden_size) - final cell state out, (hn, cn) = self.lstm(x, (h0, c0)) # We take the output from the last time step (out[:, -1, :]) # and pass it through the fully connected layer. # This is common for sequence-to-one tasks (e.g., classification). # For sequence-to-sequence, you might use 'out' directly. out = self.fc(out[:, -1, :]) return out # Example Usage: if __name__ == "__main__": input_size = 10 # e.g., 10 features per time step hidden_size = 20 # Size of the hidden state output_size = 3 # e.g., 3 classes for classification num_layers = 2 # Number of stacked LSTM layers sequence_length = 5 # Length of input sequence batch_size = 4 # Number of sequences in a batch # Create a dummy input tensor # (batch_size, sequence_length, input_size) dummy_input = torch.randn(batch_size, sequence_length, input_size) # Instantiate the LSTM model model = BasicLSTM(input_size, hidden_size, output_size, num_layers) # Perform a forward pass output = model(dummy_input) print(f"Input shape: {dummy_input.shape}") print(f"Output shape: {output.shape}") # Expected: (batch_size, output_size) # Example of a sequence-to-sequence task (e.g., predicting next sequence element) # If you wanted sequence-to-sequence, you'd modify the forward method # and likely use a TimeDistributed-like layer or process 'out' directly. # For instance, if output_size was input_size, you could aim to predict # the next element in the sequence. # A simple example without changing the FC layer: # model_seq_to_seq = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True) # out_seq_to_seq, _ = model_seq_to_seq(dummy_input) # print(f"Raw LSTM output (seq-to-seq): {out_seq_to_seq.shape}") # (batch_size, sequence_length, hidden_size) ```