Help us prove the theory of compute efficient frontier, we have to test multiple scales of the MLP in this code and plot the training results in logarithmic plot, configure multiple sizes of the same model trained on the same dataset, we have to at least have 8 configurations, add notes and considerations for what you think might me better to add for this experiment: import json
import joblib
import random

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split

# --- Set Random Seed for Reproducibility ---
def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed(42)

# --- Load Data from CSV ---
def load_data(csv_file):
    return pd.read_csv(csv_file)

# --- Data Preprocessing ---
def preprocess_data(df):
    numeric_features = [
        "Num Shaders",
        "Clock Speed (GHz)",
        "ROPs",
        "TMUs",
        "L2 Cache (MB)",
        "VRAM (GB)",
        "Bandwidth (GB/s)",
        "TFlops",
    ]
    categorical_features = ["Architecture"]
    ARCHITECTURE_MAPPING = {
        "PASCAL": 0,
        "TURING": 1,
        "AMPERE": 2,
        "ADA": 3,
        "GCN5": 4,
        "RDNA": 5,
        "RDNA2": 6,
        "RDNA3": 7,
    }

    df["Architecture"] = df["Architecture"].map(ARCHITECTURE_MAPPING).astype(int)
    target = "Relative Performance"

    # Standardize numerical features
    scaler = StandardScaler()
    X_num = scaler.fit_transform(df[numeric_features])
    X_cat = df[categorical_features].values

    y = df[target].values

    # Convert to tensors
    X_num = torch.tensor(X_num, dtype=torch.float32)
    X_cat = torch.tensor(X_cat, dtype=torch.long).squeeze()
    y = torch.tensor(y, dtype=torch.float32).unsqueeze(1)  # Unsqueeze for single target

    return X_num, X_cat, y, scaler

# --- Custom Dataset ---
class GPUDataset(Dataset):
    def __init__(self, X_num, X_cat, y):
        self.X_num = X_num
        self.X_cat = X_cat
        self.y = y

    def __len__(self):
        return len(self.y)

    def __getitem__(self, idx):
        return self.X_num[idx], self.X_cat[idx], self.y[idx]

# --- Model Architecture ---
class GPUPerformanceModel(nn.Module):
    def __init__(self, input_size, num_architectures):
        super(GPUPerformanceModel, self).__init__()
        self.arch_embedding = nn.Embedding(num_architectures, embedding_dim=4)
        self.model = nn.Sequential(
            nn.Linear(input_size + 4, 128),
            nn.ReLU(),
            nn.Dropout(0.125),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(0.125),
            nn.Linear(64, 1)
        )

    def forward(self, x_num, x_cat):
        x_cat = self.arch_embedding(x_cat)
        x = torch.cat([x_num, x_cat], dim=1)
        return self.model(x)

# --- Early Stopping ---
class EarlyStopping:
    def __init__(self, patience=50, min_delta=0):
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_loss = None
        self.early_stop = False

    def call(self, val_loss):
        if self.best_loss is None:
            self.best_loss = val_loss
        elif val_loss > self.best_loss - self.min_delta:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_loss = val_loss
            self.counter = 0

# --- Model Training ---
def train_model(model, train_loader, val_loader, epochs, learning_rate, device):
    criterion = nn.MSELoss()
    mae_loss = nn.L1Loss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=20)

    early_stopping = EarlyStopping(patience=50)

    train_losses = []
    val_losses = []
    train_mae = []
    val_mae = []
    train_rmse = []
    val_rmse = []

    for epoch in range(epochs):
        # Training
        model.train()
        epoch_train_loss = 0
        epoch_train_mae = 0
        epoch_train_rmse = 0
        for X_num_batch, X_cat_batch, y_batch in train_loader:
            X_num_batch = X_num_batch.to(device)
            X_cat_batch = X_cat_batch.to(device)
            y_batch = y_batch.to(device)

            optimizer.zero_grad()
            outputs = model(X_num_batch, X_cat_batch)
            loss = criterion(outputs, y_batch)
            loss.backward()
            optimizer.step()

            epoch_train_loss += loss.item() * X_num_batch.size(0)
            epoch_train_mae += mae_loss(outputs, y_batch).item() * X_num_batch.size(0)
            epoch_train_rmse += torch.sqrt(criterion(outputs, y_batch)).item() * X_num_batch.size(0)

        train_losses.append(epoch_train_loss / len(train_loader.dataset))
        train_mae.append(epoch_train_mae / len(train_loader.dataset))
        train_rmse.append(epoch_train_rmse / len(train_loader.dataset))

        # Validation
        model.eval()
        epoch_val_loss = 0
        epoch_val_mae = 0
        epoch_val_rmse = 0
        with torch.no_grad():
            for X_num_batch, X_cat_batch, y_batch in val_loader:
                X_num_batch = X_num_batch.to(device)
                X_cat_batch = X_cat_batch.to(device)
                y_batch = y_batch.to(device)

                outputs = model(X_num_batch, X_cat_batch)
                loss = criterion(outputs, y_batch)

                epoch_val_loss += loss.item() * X_num_batch.size(0)
                epoch_val_mae += mae_loss(outputs, y_batch).item() * X_num_batch.size(0)
                epoch_val_rmse += torch.sqrt(criterion(outputs, y_batch)).item() * X_num_batch.size(0)

        val_losses.append(epoch_val_loss / len(val_loader.dataset))
        val_mae.append(epoch_val_mae / len(val_loader.dataset))
        val_rmse.append(epoch_val_rmse / len(val_loader.dataset))

        scheduler.step(val_losses[-1])

        # Print epoch statistics
        if (epoch + 1) % 50 == 0:
            print(
                f"Epoch [{epoch+1}/{epochs}], "
                f"Train Loss: {train_losses[-1]:.4f}, Val Loss: {val_losses[-1]:.4f}, "
                f"Train MAE: {train_mae[-1]:.4f}, Val MAE: {val_mae[-1]:.4f}, "
                f"Train RMSE: {train_rmse[-1]:.4f}, Val RMSE: {val_rmse[-1]:.4f}"
            )

        # Early stopping
        #early_stopping(val_losses[-1])
        #if early_stopping.early_stop:
        #    print("Early stopping")
        #    break

    return train_losses, val_losses, train_mae, val_mae, train_rmse, val_rmse

# --- Model Evaluation ---
def evaluate_model(model, data_loader, device):
    model.eval()
    criterion = nn.MSELoss()
    mae_loss = nn.L1Loss()

    total_loss = 0
    total_mae = 0
    total_rmse = 0
    with torch.no_grad():
        for X_num_batch, X_cat_batch, y_batch in data_loader:
            X_num_batch = X_num_batch.to(device)
            X_cat_batch = X_cat_batch.to(device)
            y_batch = y_batch.to(device)

            outputs = model(X_num_batch, X_cat_batch)
            loss = criterion(outputs, y_batch)
            total_loss += loss.item() * X_num_batch.size(0)
            total_mae += mae_loss(outputs, y_batch).item() * X_num_batch.size(0)
            total_rmse += torch.sqrt(criterion(outputs, y_batch)).item() * X_num_batch.size(0)

    avg_loss = total_loss / len(data_loader.dataset)
    avg_mae = total_mae / len(data_loader.dataset)
    avg_rmse = total_rmse / len(data_loader.dataset)
    return avg_loss, avg_mae, avg_rmse

# --- Main Function ---
def main():
    # Device configuration
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Load data
    df = load_data("gpu_data.csv")

    # Preprocess data
    X_num, X_cat, y, scaler = preprocess_data(df)

    # Split data into training and validation sets
    X_num_train, X_num_val, X_cat_train, X_cat_val, y_train, y_val = train_test_split(
        X_num, X_cat, y, test_size=0.1, random_state=42
    )

    # Create datasets and loaders
    train_dataset = GPUDataset(X_num_train, X_cat_train, y_train)
    val_dataset = GPUDataset(X_num_val, X_cat_val, y_val)
    batch_size = 16
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

    # Model Parameters
    input_size = X_num.shape[1]
    num_architectures = df['Architecture'].nunique()
    # Training Parameters
    epochs = 1000
    learning_rate = 0.0011

    # Initialize model
    model = GPUPerformanceModel(
        input_size=input_size,
        num_architectures=num_architectures
    ).to(device)

    # Train model
    train_losses, val_losses, train_mae, val_mae, train_rmse, val_rmse = train_model(
        model, train_loader, val_loader, epochs=epochs, learning_rate=learning_rate, device=device
    )

    # Evaluate model on validation set
    val_dataset_full = GPUDataset(X_num_val, X_cat_val, y_val)
    val_loader_full = DataLoader(val_dataset_full, batch_size=len(val_dataset_full), shuffle=False)
    loss, mae, rmse = evaluate_model(model, val_loader_full, device)
    print(f"Validation Loss: {loss:.4f}, Validation MAE: {mae:.4f}, Validation RMSE: {rmse:.4f}")

    # Save model and scaler
    model_name = "GPU-PERF-Improved"
    torch.save(model.state_dict(), f"{model_name}.pth")
    joblib.dump(scaler, f"{model_name}_scaler.pkl")

    # Save model configuration
    model_config = {
        "model_params": {
            "input_size": input_size,
            "num_architectures": num_architectures,
            "embedding_dim": 4,
            "layers": [128, 64, 1],
            "dropout": 0.125,
        },
        "training_params": {
            "epochs": epochs,
            "optimizer": "Adam",
            "learning_rate": learning_rate,
            "loss_function": "MSELoss",
            "batch_size": batch_size,
            "early_stopping_patience": 50,
            "scheduler": "ReduceLROnPlateau",
        },
        "validation_loss": loss,
        "validation_mae": mae,
        "validation_rmse": rmse,
    }
    with open(f"{model_name}.json", "w") as f:
        json.dump(model_config, f)

    # Plot training and validation loss curves
    plt.figure(figsize=(10, 5))
    plt.plot(train_losses, label="Training Loss")
    plt.plot(val_losses, label="Validation Loss")
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend()
    plt.title("Training and Validation Loss Curves")
    plt.savefig(f"{model_name}_loss_curve.png")
    plt.show()

    # Plot training and validation MAE curves
    plt.figure(figsize=(10, 5))
    plt.plot(train_mae, label="Training MAE")
    plt.plot(val_mae, label="Validation MAE")
    plt.xlabel("Epochs")
    plt.ylabel("MAE")
    plt.legend()
    plt.title("Training and Validation MAE Curves")
    plt.savefig(f"{model_name}_mae_curve.png")
    plt.show()

    # Plot predicted vs actual performance
    model.eval()
    with torch.no_grad():
        X_num_val = X_num_val.to(device)
        X_cat_val = X_cat_val.to(device)
        y_val_pred = model(X_num_val, X_cat_val).cpu().numpy()
    y_val_true = y_val.cpu().numpy()

    plt.figure(figsize=(10, 5))
    plt.scatter(y_val_true, y_val_pred, color="blue", label="Predictions", marker='x')
    plt.xlabel("Actual Performance")
    plt.ylabel("Predicted Performance")
    plt.legend()
    plt.title("Predicted vs Actual GPU Performance")
    plt.savefig(f"{model_name}_pred_vs_actual.png")
    plt.show()

if __name__ == "__main__":
    main()