Edits history of script submission #6061 for ' Train RNN for time series prediction (prediction)'

  • python3
    # Import necessary libraries
    import torch
    import torch.nn as nn
    from typing import List
    
    
    # Define a simple RNN model for time series prediction
    class RNNModel(nn.Module):
        def __init__(
            self, input_size: int, hidden_size: int, output_size: int, num_layers: int
        ):
            super(RNNModel, self).__init__()
            self.hidden_size = hidden_size
            self.num_layers = num_layers
            self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True)
            self.fc = nn.Linear(hidden_size, output_size)
    
        def forward(self, x):
            # Initialize hidden and cell states
            h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size)
            # Forward propagate RNN
            out, _ = self.rnn(x, h0)
            # Pass the output of the last time step to the classifier
            out = self.fc(out[:, -1, :])
            return out
    
    
    def main(
        data: List[float], num_epochs: int = 100, learning_rate: float = 0.01
    ) -> List[float]:
        """
        Perform time series prediction using an RNN model.
    
        Parameters:
        - data: List[float], the time series data for training.
        - num_epochs: int, the number of epochs to train the model.
        - learning_rate: float, the learning rate for the optimizer.
    
        Returns:
        - predictions: List[float], the predicted values for the time series.
        """
        # Convert data to PyTorch tensors
        data_normalized = torch.FloatTensor(data).view(-1)
        # Define the model
        input_size = 1
        hidden_size = 64
        output_size = 1
        num_layers = 1
        model = RNNModel(input_size, hidden_size, output_size, num_layers)
        # Loss and optimizer
        criterion = nn.MSELoss()
        optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    
        # Train the model
        for epoch in range(num_epochs):
            for i in range(len(data_normalized) - 1):
                # Prepare data
                sequence = data_normalized[i : i + 1].view(-1, 1, 1)
                target = data_normalized[i + 1].view(-1)
                # Forward pass
                output = model(sequence)
                loss = criterion(output.view(-1), target)
                # Backward and optimize
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
    
            if (epoch + 1) % 10 == 0:
                print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")
    
        # Predict (Here we use the last part of the data as a simple example)
        test_data = data_normalized[-1:].view(-1, 1, 1)
        with torch.no_grad():
            predictions = model(test_data).view(-1).tolist()
    
        return predictions
    

    Submitted by henri186 763 days ago