recurrent_nn.py

import torch
import torch.nn as nn

class simple_nn(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super().__init__()
        self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True)
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fnn = nn.Linear(hidden_size, output_size)

    def forward(self, x): # x: [batch_size, seq_len, input_size]
    
        #out, _ = self.rnn(x) # out: [batch_size, seq_len, hidden_size]
        out, _ = self.lstm(x) # out: [batch_size, seq_len, hidden_size]
        out = self.fnn(out[:, -1, :]) # Use the output of the final time-step for prediction.
        return out


batch_size = 32
seq_len = 10
input_size = 8
hidden_size = 20
num_layers = 2
output_size = 1

model = simple_nn(input_size, hidden_size, num_layers, output_size)
x = torch.randn(batch_size, seq_len, input_size)
y = model(x)
print(y.shape)  # Expected shape: [32, 1]

# training process is same as transformer file.