mlp_tabular_regression.py

# -*- coding: utf-8 -*-
"""MLP_Tabular_Regression.ipynb
"""

import openml
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
from torch.nn.utils import clip_grad_norm_
from tqdm import tqdm
import os

# Choose device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ============================
# 2. Download and Prepare Dataset from OpenML (ID 44140)
# ============================
# Download the dataset from OpenML

dataset_id = 507
dataset_name = str(dataset_id)
dataset = openml.datasets.get_dataset(dataset_id)
print("Dataset name:", dataset.name)

# Get the dataset as a pandas DataFrame
df, *_ = dataset.get_data(dataset_format="dataframe")
target_attr = dataset.default_target_attribute
print("Target attribute:", target_attr)

# Separate features and target
X = df.drop(columns=[target_attr])
y = df[target_attr]

# Ensure the target is numeric for regression
y = y.astype(float)

# Fill missing values in features and target (if any)
X.fillna(X.mean(), inplace=True)
y.fillna(y.mean(), inplace=True)

# Standardize the features using StandardScaler
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# Split into training and test sets (80/20 split)
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)
print("Train shape:", X_train.shape, y_train.shape)
print("Test shape:", X_test.shape, y_test.shape)

# ============================
# 3. Convert Data to PyTorch Tensors and Create DataLoaders
# ============================
X_train_tensor = torch.from_numpy(X_train).float()
X_test_tensor  = torch.from_numpy(X_test).float()

# For regression, we want y to be of shape (N, 1)
y_train_tensor = torch.from_numpy(y_train.values if isinstance(y_train, pd.Series) else y_train).float().unsqueeze(1)
y_test_tensor  = torch.from_numpy(y_test.values if isinstance(y_test, pd.Series) else y_test).float().unsqueeze(1)

train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
test_dataset  = TensorDataset(X_test_tensor, y_test_tensor)

batch_size = 64
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
test_loader  = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)

# ============================
# 4. Define the MLP Model for Regression
# ============================

# Define the MLP Model for Regression with Dropout
class MLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, output_dim=1, dropout=0.2):
        """
        Args:
            input_dim: Number of input features.
            hidden_dim: Number of neurons in each hidden layer.
            num_layers: Number of hidden layers.
            output_dim: Number of outputs (1 for regression).
            dropout: Dropout probability to use after each hidden layer.
        """
        super(MLP, self).__init__()
        layers = []
        # First hidden layer: input_dim -> hidden_dim, then ReLU and dropout.
        layers.append(nn.Linear(input_dim, hidden_dim))
        layers.append(nn.ReLU())
        layers.append(nn.Dropout(dropout))

        # Additional hidden layers
        for _ in range(num_layers - 1):
            layers.append(nn.Linear(hidden_dim, hidden_dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout))

        # Final output layer: hidden_dim -> output_dim (no activation or dropout)
        layers.append(nn.Linear(hidden_dim, output_dim))
        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)

# Set MLP parameters (adjust as needed)
input_dim = X_train.shape[1]  # assuming X_train has been defined earlier
hidden_dim = 64               # N neurons per hidden layer
num_layers = 5                # L hidden layers
output_dim = 1                # For regression

def create_mlp():
    return MLP(input_dim, hidden_dim, num_layers, output_dim, dropout=0.2)

# Modified FastAdam optimizer with logging of Eve updates.
class FastAdam:
    def __init__(self, params, lr=0.001, betas=(0.9, 0.999), eps=1e-8,
                 custom=False, weight_decay=1e-4, c=1.0,
                 log_csv="/content/gdrive/My Drive/Colab Notebooks/Eve/eve_updates_log.csv"):
        self.lr = lr
        self.beta1, self.beta2 = betas
        self.eps = eps
        self.custom = custom
        self.weight_decay = weight_decay
        self.t = 0
        self.c = c

        # Open CSV file and write header for aggregated logging.
        if log_csv is not None:
            self.csv_file = open(log_csv, 'w', newline='')
            self.csv_writer = csv.writer(self.csv_file)
            self.csv_writer.writerow(["step", "total_eve_updates", "layer_counts"])
        else:
            self.csv_file = None
            self.csv_writer = None

        # Separate names and tensors.
        self.names = []
        self.params = []
        for i, p in enumerate(params):
            if isinstance(p, tuple):
                self.names.append(p[0])
                self.params.append(p[1])
            else:
                self.names.append("param_" + str(i))
                self.params.append(p)

        # Initialize first and second moment estimates and previous gradients.
        self.m = [torch.zeros_like(p) for p in self.params]
        self.v = [torch.zeros_like(p) for p in self.params]
        self.prev_grad = [torch.zeros_like(p) for p in self.params]

    def zero_grad(self):
        for p in self.params:
            if p.grad is not None:
                p.grad.zero_()

    def step(self):
        self.t += 1

        # Initialize aggregation counters.
        total_eve_update_count = 0
        layer_counts = {}

        for i, (p, m, v, prev_g) in enumerate(zip(self.params, self.m, self.v, self.prev_grad)):
            if p.grad is None:
                continue
            grad = p.grad

            # Apply decoupled weight decay (AdamW style)
            if self.weight_decay != 0:
                p.data.mul_(1 - self.lr * self.weight_decay)

            # Update biased first and second moment estimates.
            m[:] = self.beta1 * m + (1 - self.beta1) * grad
            v[:] = self.beta2 * v + (1 - self.beta2) * grad.square()

            # Compute bias-corrected estimates.
            m_hat = m / (1 - self.beta1 ** self.t)
            v_hat = v / (1 - self.beta2 ** self.t)

            # Compute standard Adam update.
            update = -self.lr * m_hat / (v_hat.sqrt() + self.eps)

            if self.custom:
                # Condition: use Adam update if its magnitude is less than the previous gradient's magnitude.
                condition = update.abs() <= self.c * prev_g.abs()
                # backward_mask is True where the backward update (Eve correction) is applied.
                backward_mask = ~condition
                count = int(backward_mask.sum().item())
                total_eve_update_count += count

                # Extract layer name (using the part before the first dot).
                layer = self.names[i].split('.')[0] if '.' in self.names[i] else self.names[i]
                layer_counts[layer] = layer_counts.get(layer, 0) + count

                # Where condition is False, apply the backward update (using 0.1 * previous gradient).
                update = torch.where(condition, update, -0.1 * prev_g)

            # Update the parameter.
            p.data.add_(update)
            # Store the current gradient for the next iteration.
            prev_g[:] = grad

        # Log aggregated Eve update information for this step.
        if self.csv_writer is not None:
            self.csv_writer.writerow([self.t, total_eve_update_count, str(layer_counts)])

    def close(self):
        if self.csv_file is not None:
            self.csv_file.close()

    def __del__(self):
        self.close()

# ============================
# 6. Define run_experiment Function for Regression
# ============================
def run_experiment(optimizer_name, model_fn, train_loader, test_loader,
                   num_epochs=10, lr=1e-3, momentum=0.9, max_norm=1.0):
    """
    - Creates the model using model_fn().
    - Chooses the optimizer based on optimizer_name:
         'fastadam', 'adam', 'sgd', 'adam_gc' (Adam with gradient clipping).
    - For each epoch, trains the model on the training data and evaluates on test data.
    - The loss is MSELoss for regression.
    - Returns (model, history).
    Additionally, displays progress bars for training and validation.
    """
    model = model_fn().to(device)

    do_grad_clip = False

    if optimizer_name.lower() == 'fastadam':
        optimizer = FastAdam(model.parameters(), lr=lr, custom=True, weight_decay=1e-4)
    elif optimizer_name.lower() == 'adam':
        optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
    elif optimizer_name.lower() == 'sgd':
        optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum, weight_decay=1e-4)
    elif optimizer_name.lower() == 'adam_gc':
        optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
        do_grad_clip = True
    else:
        raise ValueError(f"Optimizer '{optimizer_name}' not recognized.")

    criterion = nn.MSELoss()
    history = {
        'train_loss': [],
        'test_loss': []
    }

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        total = 0

        train_pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs} Training", leave=False)
        for X_batch, y_batch in train_pbar:
            X_batch = X_batch.to(device)
            y_batch = y_batch.to(device)

            optimizer.zero_grad()
            outputs = model(X_batch)
            loss = criterion(outputs, y_batch)
            loss.backward()

            if do_grad_clip:
                clip_grad_norm_(model.parameters(), max_norm)

            optimizer.step()

            running_loss += loss.item() * X_batch.size(0)
            total += X_batch.size(0)
            train_pbar.set_postfix(loss=f"{loss.item():.4f}")

        tr_loss = running_loss / total

        # Evaluation on test set
        model.eval()
        running_loss_val = 0.0
        total_val = 0
        val_pbar = tqdm(test_loader, desc=f"Epoch {epoch+1}/{num_epochs} Validation", leave=False)
        with torch.no_grad():
            for X_batch, y_batch in val_pbar:
                X_batch = X_batch.to(device)
                y_batch = y_batch.to(device)
                outputs = model(X_batch)
                loss_val = criterion(outputs, y_batch)
                running_loss_val += loss_val.item() * X_batch.size(0)
                total_val += X_batch.size(0)
                val_pbar.set_postfix(loss=f"{loss_val.item():.4f}")
        val_loss = running_loss_val / total_val

        history['train_loss'].append(tr_loss)
        history['test_loss'].append(val_loss)
        print(f"[{optimizer_name.upper()}] Epoch [{epoch+1}/{num_epochs}] - Train MSE: {tr_loss:.4f}, Test MSE: {val_loss:.4f}")

    return model, history

save_dir = "/content/gdrive/My Drive/Colab Notebooks/Eve/mlp_regression/"

optimizers_list = ['fastadam',
                   #'adam',
                   #'sgd',
                   #'adam_gc'
                   ]
results = {}
final_models = {}

num_epochs = 400
learning_rate = 1e-3
momentum = 0.9
max_norm = 1.0

for opt_name in optimizers_list:
    print(f"\n*** Training with {opt_name.upper()} ***")
    model, history = run_experiment(
        opt_name,
        create_mlp,  # our MLP creation function
        train_loader,
        test_loader,
        num_epochs=num_epochs,
        lr=learning_rate,
        momentum=momentum,
        max_norm=max_norm
    )
    results[opt_name] = history
    final_models[opt_name] = model

    # Salva il modello (solo lo state_dict per maggiore compatibilità)
    model_save_path = os.path.join(save_dir, f"{opt_name}_{dataset_name}_{str(hidden_dim)}_{str(num_layers)}_model_state_dict_-.pth")
    torch.save(model.state_dict(), model_save_path)
    print(f"Modello {opt_name} salvato in: {model_save_path}")