Create app.py

app.py

@@ -0,0 +1,479 @@
+import io
+import random
+import tempfile
+from dataclasses import dataclass
+
+import gradio as gr
+import matplotlib
+
+matplotlib.use("Agg")  # headless-friendly for Hugging Face Spaces
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, TensorDataset
+
+import lightning.pytorch as pl
+
+
+# -----------------------------
+# Data
+# -----------------------------
+@dataclass
+class DataSpec:
+    n_samples: int = 1024
+    n_features: int = 10
+    noise_std: float = 0.3
+    train_frac: float = 0.8
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def make_synthetic_regression(spec: DataSpec, seed: int = 42):
+    """
+    Synthetic regression:
+      y = X @ w_true + b_true + noise
+
+    X: (n_samples, 10)
+    y: (n_samples, 1)
+    """
+    set_seed(seed)
+
+    w_true = torch.randn(spec.n_features, 1) * 2.0
+    b_true = torch.randn(1) * 0.5
+
+    X = torch.randn(spec.n_samples, spec.n_features)
+    noise = torch.randn(spec.n_samples, 1) * spec.noise_std
+    y = X @ w_true + b_true + noise
+
+    n_train = int(spec.n_samples * spec.train_frac)
+    X_train, y_train = X[:n_train], y[:n_train]
+    X_val, y_val = X[n_train:], y[n_train:]
+
+    return X_train, y_train, X_val, y_val, w_true, b_true
+
+
+def build_full_dataset_df(X_train, y_train, X_val, y_val) -> pd.DataFrame:
+    cols = [f"x{i}" for i in range(10)]
+
+    train_df = pd.DataFrame(X_train.cpu().numpy(), columns=cols)
+    train_df["y"] = y_train.cpu().numpy().reshape(-1)
+    train_df["split"] = "train"
+
+    val_df = pd.DataFrame(X_val.cpu().numpy(), columns=cols)
+    val_df["y"] = y_val.cpu().numpy().reshape(-1)
+    val_df["split"] = "val"
+
+    return pd.concat([train_df, val_df], ignore_index=True)
+
+
+def save_df_to_temp_csv(df: pd.DataFrame) -> str:
+    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv", prefix="synthetic_regression_")
+    df.to_csv(tmp.name, index=False)
+    return tmp.name
+
+
+# -----------------------------
+# Plot helper
+# -----------------------------
+def fig_to_image(fig) -> np.ndarray:
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png", bbox_inches="tight", dpi=160)
+    plt.close(fig)
+    buf.seek(0)
+    return plt.imread(buf)
+
+
+def plot_losses(train_losses, val_losses, title: str) -> np.ndarray:
+    fig = plt.figure()
+    if len(train_losses) > 0:
+        plt.plot(range(1, len(train_losses) + 1), train_losses, marker="o", label="train")
+    if len(val_losses) > 0:
+        plt.plot(range(1, len(val_losses) + 1), val_losses, marker="o", label="val")
+    plt.xlabel("Epoch")
+    plt.ylabel("MSE Loss")
+    plt.title(title)
+    plt.grid(True, alpha=0.3)
+    plt.legend()
+    return fig_to_image(fig)
+
+
+def weights_table(w_true: torch.Tensor, w_learned: torch.Tensor) -> pd.DataFrame:
+    rows = []
+    for i in range(10):
+        wt = float(w_true[i].item())
+        wl = float(w_learned[i].item())
+        rows.append(
+            {"feature": f"x{i}", "w_true": wt, "w_learned": wl, "abs_error": abs(wt - wl)}
+        )
+    df = pd.DataFrame(rows).round(4)
+    df = df.sort_values("abs_error", ascending=False).reset_index(drop=True)
+    return df
+
+
+# -----------------------------
+# Raw PyTorch training
+# -----------------------------
+def train_raw(
+    X_train, y_train, X_val, y_val,
+    init_state_dict,
+    lr: float,
+    batch_size: int,
+    epochs: int,
+    device: torch.device
+):
+    train_loader = DataLoader(
+        TensorDataset(X_train, y_train),
+        batch_size=batch_size,
+        shuffle=False,  # fixed order to make raw vs lightning comparable
+        num_workers=0,
+    )
+    val_loader = DataLoader(
+        TensorDataset(X_val, y_val),
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=0,
+    )
+
+    model = nn.Linear(10, 1)
+    model.load_state_dict(init_state_dict)
+    model.to(device)
+
+    loss_fn = nn.MSELoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
+
+    train_losses, val_losses = [], []
+
+    for _epoch in range(epochs):
+        model.train()
+        running, seen = 0.0, 0
+
+        for x, y in train_loader:
+            x, y = x.to(device), y.to(device)
+
+            optimizer.zero_grad()
+            y_pred = model(x)
+            loss = loss_fn(y_pred, y)
+            loss.backward()
+            optimizer.step()
+
+            bs = x.size(0)
+            running += loss.item() * bs
+            seen += bs
+
+        train_losses.append(running / max(seen, 1))
+
+        model.eval()
+        running, seen = 0.0, 0
+        with torch.no_grad():
+            for x, y in val_loader:
+                x, y = x.to(device), y.to(device)
+                y_pred = model(x)
+                loss = loss_fn(y_pred, y)
+                bs = x.size(0)
+                running += loss.item() * bs
+                seen += bs
+
+        val_losses.append(running / max(seen, 1))
+
+    with torch.no_grad():
+        w_learned = model.weight.detach().view(-1, 1).cpu()
+        b_learned = model.bias.detach().view(1).cpu()
+
+    return train_losses, val_losses, w_learned, b_learned
+
+
+# -----------------------------
+# Lightning training
+# -----------------------------
+class LitModel(pl.LightningModule):
+    def __init__(self, lr: float, init_state_dict):
+        super().__init__()
+        self.save_hyperparameters(ignore=["init_state_dict"])
+        self.model = nn.Linear(10, 1)
+        self.model.load_state_dict(init_state_dict)
+        self.lr = lr
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, _batch_idx):
+        x, y = batch
+        loss = F.mse_loss(self(x), y)
+        self.log("train_loss", loss, on_step=False, on_epoch=True)
+        return loss
+
+    def validation_step(self, batch, _batch_idx):
+        x, y = batch
+        loss = F.mse_loss(self(x), y)
+        self.log("val_loss", loss, on_step=False, on_epoch=True)
+        return loss
+
+    def configure_optimizers(self):
+        return torch.optim.SGD(self.parameters(), lr=self.lr)
+
+
+class LossHistoryCallback(pl.Callback):
+    def __init__(self):
+        self.train_losses = []
+        self.val_losses = []
+
+    def on_train_epoch_end(self, trainer, pl_module):
+        m = trainer.callback_metrics
+        if "train_loss" in m:
+            self.train_losses.append(float(m["train_loss"].detach().cpu().item()))
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        m = trainer.callback_metrics
+        if "val_loss" in m:
+            self.val_losses.append(float(m["val_loss"].detach().cpu().item()))
+
+
+def train_lightning(
+    X_train, y_train, X_val, y_val,
+    init_state_dict,
+    lr: float,
+    batch_size: int,
+    epochs: int,
+    device_choice: str,
+    seed: int
+):
+    pl.seed_everything(seed, workers=True)
+
+    train_loader = DataLoader(
+        TensorDataset(X_train, y_train),
+        batch_size=batch_size,
+        shuffle=False,  # fixed order to make raw vs lightning comparable
+        num_workers=0,
+    )
+    val_loader = DataLoader(
+        TensorDataset(X_val, y_val),
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=0,
+    )
+
+    want_cuda = (device_choice == "cuda")
+    has_cuda = torch.cuda.is_available()
+    using_cuda = want_cuda and has_cuda
+    accelerator = "gpu" if using_cuda else "cpu"
+
+    model = LitModel(lr=lr, init_state_dict=init_state_dict)
+    history = LossHistoryCallback()
+
+    trainer = pl.Trainer(
+        max_epochs=epochs,
+        accelerator=accelerator,
+        devices=1,
+        deterministic=True,
+        logger=False,
+        enable_checkpointing=False,
+        enable_progress_bar=False,
+        enable_model_summary=False,
+        callbacks=[history],
+    )
+
+    trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)
+
+    with torch.no_grad():
+        w_learned = model.model.weight.detach().view(-1, 1).cpu()
+        b_learned = model.model.bias.detach().view(1).cpu()
+
+    return history.train_losses, history.val_losses, w_learned, b_learned, using_cuda
+
+
+# -----------------------------
+# Run BOTH on same data & same init weights
+# -----------------------------
+def run_both(
+    n_samples: int,
+    noise_std: float,
+    lr: float,
+    batch_size: int,
+    epochs: int,
+    seed: int,
+    device_choice: str,
+):
+    # 1) Generate data ONCE
+    spec = DataSpec(n_samples=n_samples, n_features=10, noise_std=noise_std, train_frac=0.8)
+    X_train, y_train, X_val, y_val, w_true, b_true = make_synthetic_regression(spec, seed=seed)
+
+    # Preview + CSV download
+    preview_n = min(20, X_train.shape[0])
+    df_preview = pd.DataFrame(
+        X_train[:preview_n].cpu().numpy(),
+        columns=[f"x{i}" for i in range(10)]
+    )
+    df_preview["y"] = y_train[:preview_n].cpu().numpy().reshape(-1)
+    df_preview = df_preview.round(4)
+
+    full_df = build_full_dataset_df(X_train, y_train, X_val, y_val).round(4)
+    csv_path = save_df_to_temp_csv(full_df)
+
+    # 2) Create ONE initial weight state and reuse it for both trainings
+    set_seed(seed + 123)  # separate seed so "data seed" vs "init seed" is clear & repeatable
+    base = nn.Linear(10, 1)
+    init_state = base.state_dict()
+
+    # 3) Train RAW
+    if device_choice == "cuda" and torch.cuda.is_available():
+        device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+
+    raw_train_losses, raw_val_losses, raw_w, raw_b = train_raw(
+        X_train, y_train, X_val, y_val,
+        init_state_dict=init_state,
+        lr=lr,
+        batch_size=batch_size,
+        epochs=epochs,
+        device=device
+    )
+
+    raw_loss_img = plot_losses(raw_train_losses, raw_val_losses, "Raw PyTorch (Manual Loop)")
+    raw_weights_df = weights_table(w_true.cpu(), raw_w)
+
+    raw_summary = (
+        f"Device: {device}\n"
+        f"Final train loss: {raw_train_losses[-1]:.6f}\n"
+        f"Final val loss:   {raw_val_losses[-1]:.6f}\n\n"
+        f"True bias (b_true):         {float(b_true.item()):.4f}\n"
+        f"Learned bias (raw):         {float(raw_b.item()):.4f}\n"
+    )
+
+    # 4) Train LIGHTNING (same data + same init weights)
+    lt_train_losses, lt_val_losses, lt_w, lt_b, using_cuda = train_lightning(
+        X_train, y_train, X_val, y_val,
+        init_state_dict=init_state,
+        lr=lr,
+        batch_size=batch_size,
+        epochs=epochs,
+        device_choice=device_choice,
+        seed=seed + 999,
+    )
+
+    lt_loss_img = plot_losses(lt_train_losses, lt_val_losses, "Lightning (Trainer.fit)")
+    lt_weights_df = weights_table(w_true.cpu(), lt_w)
+
+    lt_summary = (
+        f"Requested device: {device_choice}\n"
+        f"Using device: {'cuda' if using_cuda else 'cpu'}\n"
+        f"Final train loss: {lt_train_losses[-1]:.6f}\n"
+        f"Final val loss:   {lt_val_losses[-1]:.6f}\n\n"
+        f"True bias (b_true):         {float(b_true.item()):.4f}\n"
+        f"Learned bias (lightning):   {float(lt_b.item()):.4f}\n"
+    )
+
+    raw_snippet = """# Raw PyTorch: manual training loop
+model = nn.Linear(10, 1)
+optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+loss_fn = nn.MSELoss()
+
+for x, y in dataloader:
+    optimizer.zero_grad()
+    y_pred = model(x)
+    loss = loss_fn(y_pred, y)
+    loss.backward()
+    optimizer.step()
+"""
+
+    lightning_snippet = """# Lightning: training logic organized in a class
+import lightning.pytorch as pl
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+class LitModel(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.model = nn.Linear(10, 1)
+
+    def training_step(self, batch, _):
+        x, y = batch
+        return F.mse_loss(self.model(x), y)
+
+    def configure_optimizers(self):
+        return torch.optim.SGD(self.parameters(), lr=0.01)
+
+# pl.Trainer(max_epochs=1).fit(LitModel(), dataloader)
+"""
+
+    return (
+        df_preview,
+        csv_path,
+        raw_loss_img, raw_weights_df, raw_summary, raw_snippet,
+        lt_loss_img, lt_weights_df, lt_summary, lightning_snippet
+    )
+
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
+with gr.Blocks(title="Raw PyTorch vs Lightning (Same Data)") as demo:
+    gr.Markdown(
+        """
+# Raw PyTorch vs PyTorch Lightning — Same Data, Same Initialization
+
+This Space trains **two versions** of the same model on the **same synthetic dataset**:
+
+- **Raw PyTorch**: manual training loop (`zero_grad → forward → loss → backward → step`)
+- **Lightning**: training organized in `LightningModule` + `Trainer.fit(...)`
+
+To make comparisons fair:
+- The dataset is generated once per run using the same seed
+- The **initial model weights are copied** so both start identically
+- Batch order is fixed (no shuffle) so both see batches in the same order
+"""
+    )
+
+    with gr.Row():
+        n_samples = gr.Slider(256, 8192, value=1024, step=256, label="Number of samples")
+        noise_std = gr.Slider(0.0, 2.0, value=0.3, step=0.05, label="Noise (std dev)")
+
+    with gr.Row():
+        lr = gr.Slider(1e-4, 1.0, value=0.01, step=1e-4, label="Learning rate (SGD)")
+        batch_size = gr.Dropdown([16, 32, 64, 128, 256], value=64, label="Batch size")
+
+    with gr.Row():
+        epochs = gr.Slider(1, 50, value=10, step=1, label="Epochs")
+        seed = gr.Number(value=42, precision=0, label="Seed (controls data)")
+
+    device_choice = gr.Radio(["cpu", "cuda"], value="cpu", label="Device (cuda only if available)")
+    run_btn = gr.Button("Generate Data + Train BOTH", variant="primary")
+
+    with gr.Tab("Data"):
+        data_preview = gr.Dataframe(label="First 20 rows of TRAIN split", wrap=True)
+        download_file = gr.File(label="Download full dataset CSV (train + val)")
+
+    with gr.Tab("Raw PyTorch"):
+        raw_loss_img = gr.Image(label="Loss Curve (Raw)", type="numpy")
+        raw_weights_df = gr.Dataframe(label="Weights: True vs Learned (Raw)", wrap=True)
+        raw_summary_txt = gr.Textbox(label="Summary (Raw)", lines=8)
+        raw_code = gr.Code(label="Raw loop snippet", language="python")
+
+    with gr.Tab("Lightning"):
+        lt_loss_img = gr.Image(label="Loss Curve (Lightning)", type="numpy")
+        lt_weights_df = gr.Dataframe(label="Weights: True vs Learned (Lightning)", wrap=True)
+        lt_summary_txt = gr.Textbox(label="Summary (Lightning)", lines=8)
+        lt_code = gr.Code(label="Lightning snippet", language="python")
+
+    run_btn.click(
+        fn=run_both,
+        inputs=[n_samples, noise_std, lr, batch_size, epochs, seed, device_choice],
+        outputs=[
+            data_preview,
+            download_file,
+            raw_loss_img, raw_weights_df, raw_summary_txt, raw_code,
+            lt_loss_img, lt_weights_df, lt_summary_txt, lt_code,
+        ],
+    )
+
+if __name__ == "__main__":
+    demo.launch()