Switch from TensorFlow (Keras) to PyTorch

[cbrnr]

Since TensorFlow (Keras) is lagging behind current Python versions (they don't have wheels for 3.13 yet), we should consider switching to PyTorch, which seems to be much more actively developed.

This means we'd have to adapt our model computation script in examples/classifiers, I'm not sure how much work this is (but it's probably worth it).

[cbrnr]

As a starting point, here's what o1 suggested:

# ...existing imports...
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader

# ...existing code...

# Define a PyTorch model similar to the original Keras model
class WRNGRUMesa(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, mask_value=-1.0):
        super().__init__()
        # We'll use LayerNorm instead of BatchNorm1d for simplicity with sequences
        self.layer_norm = nn.LayerNorm(input_dim)
        self.fc = nn.Linear(input_dim, 64)
        self.gru1 = nn.GRU(64, hidden_dim, batch_first=True, bidirectional=True)
        self.gru2 = nn.GRU(hidden_dim * 2, hidden_dim, batch_first=True, bidirectional=True)
        self.output = nn.Linear(hidden_dim * 2, output_dim)
        self.mask_value = mask_value

    def forward(self, x):
        # x shape: [batch_size, seq_len, input_dim]
        # Replace masked values with 0 to reduce their impact
        x = torch.where(x == self.mask_value, torch.zeros_like(x), x)

        # Apply layer normalization across the input dimension
        x = self.layer_norm(x)

        # Dense layer + ReLU
        x = F.relu(self.fc(x))

        # Two bidirectional GRU layers
        x, _ = self.gru1(x)  # shape: [batch_size, seq_len, hidden_dim*2]
        x, _ = self.gru2(x)  # shape: [batch_size, seq_len, hidden_dim*2]

        # Final linear layer (logits)
        x = self.output(x)   # shape: [batch_size, seq_len, output_dim]
        return x

if TRAIN:
    print("‣  Starting training (PyTorch)...")

    # Convert your numpy arrays to torch tensors
    X_train = torch.from_numpy(features_train_pad).float()  # [N, seq_len, num_features]
    # Keras used one-hot labels; convert them back to class indices for PyTorch
    y_train = torch.argmax(torch.from_numpy(stages_train_pad), dim=-1).long()  # [N, seq_len]

    train_dataset = TensorDataset(X_train, y_train)
    train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)

    # Create model, loss function, and optimizer
    model = WRNGRUMesa(
        input_dim=features_train_pad.shape[2],
        hidden_dim=8,
        output_dim=stages_train_pad.shape[-1],
        mask_value=-1.0,
    )
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.RMSprop(model.parameters(), lr=1e-3)

    model.train()
    for epoch in range(25):
        epoch_loss = 0.0
        for batch_x, batch_y in train_loader:
            optimizer.zero_grad()
            outputs = model(batch_x)  # [batch_size, seq_len, output_dim]

            # Reshape for cross-entropy: [batch_size*seq_len, output_dim], [batch_size*seq_len]
            N, T, C = outputs.size()
            loss = criterion(outputs.view(N * T, C), batch_y.view(N * T))
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()

        print(f"Epoch {epoch+1}/25, Loss: {epoch_loss/len(train_loader):.4f}")

    # ...save_classifier logic can remain the same, but you'll need to adapt it
    # to store a PyTorch model instead of a Keras model.
    # e.g., torch.save(model.state_dict(), "wrn-gru-mesa.pt")

# ...existing code for testing/evaluation (adapt for PyTorch if needed)...