Raster Example implementation about pyxlstm HOT 1 CLOSED

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from xLSTM.mlstm import mLSTM
from xLSTM.utils import config, logging, utils

# Simulate raster data
def generate_synthetic_raster_data(num_timesteps, width, height):
    return np.random.rand(num_timesteps, width, height)

# Generate synthetic raster data
num_timesteps = 10
width = 5
height = 5
raster_data = generate_synthetic_raster_data(num_timesteps, width, height)

# Normalize the raster data
raster_data = raster_data / np.max(raster_data)

# Convert raster data to PyTorch tensor
raster_data = torch.Tensor(raster_data).unsqueeze(-1)  # Add channel dimension

# Define the mLSTM model
class RasterTimeSeriesModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super(RasterTimeSeriesModel, self).__init__()
        self.mlstm = mLSTM(input_dim, hidden_dim, num_layers)
        self.fc = nn.Linear(hidden_dim, output_dim)
    
    def forward(self, x):
        h, _ = self.mlstm(x)
        out = self.fc(h)
        return out

# Model parameters
input_dim = width * height  # Flattened raster data
hidden_dim = 128
output_dim = input_dim  # Output should match the flattened input size
num_layers = 2  # Number of layers in mLSTM

# Reshape raster data to (batch_size, sequence_length, input_dim)
raster_data = raster_data.view(num_timesteps, -1).unsqueeze(0)  # Add batch dimension

model = RasterTimeSeriesModel(input_dim, hidden_dim, output_dim, num_layers)

# Training parameters
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Split data into train and test sets
train_data = raster_data[:, :8, :]  # Using first 8 timesteps for training
test_data = raster_data[:, 8:, :]  # Using the last 2 timesteps for testing

# Training loop
for epoch in range(100):  # Number of epochs
    model.train()
    optimizer.zero_grad()
    
    output = model(train_data)
    loss = criterion(output, train_data)
    loss.backward()
    optimizer.step()
    
    if epoch % 10 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item()}')

# Step 4: Evaluate the model
model.eval()
with torch.no_grad():
    output = model(test_data)
    test_loss = criterion(output, test_data)
    print(f'Test Loss: {test_loss.item()}')