{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":3649,"databundleVersionId":46718,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":15908653,"datasetId":10201741,"databundleVersionId":16864351}],"dockerImageVersionId":31329,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-04-24T05:15:00.239078Z","iopub.execute_input":"2026-04-24T05:15:00.239385Z","iopub.status.idle":"2026-04-24T05:15:01.367081Z","shell.execute_reply.started":"2026-04-24T05:15:00.239348Z","shell.execute_reply":"2026-04-24T05:15:01.366031Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"\nSelf-Pruning Neural Network on CIFAR-10\n========================================\nImplements a feed-forward network with learnable gate parameters that \nencourage sparsity via L1 regularization, enabling the network to prune\nits own weights during training.\n\nAuthor: AI Engineer Case Study\n\"\"\"\n\nimport os\nimport math\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader\nimport torchvision\nimport torchvision.transforms as transforms\nimport numpy as np\nimport matplotlib\nmatplotlib.use('Agg')  # Non-interactive backend for headless environments\nimport matplotlib.pyplot as plt\nimport json\n\n\n# ─────────────────────────────────────────────────────────────────────────────\n# Part 1: PrunableLinear Layer\n# ─────────────────────────────────────────────────────────────────────────────\n\nclass PrunableLinear(nn.Module):\n    \"\"\"\n    A custom linear layer augmented with learnable gate parameters.\n\n    Each weight w_ij has a corresponding gate_score g_ij. The gate is\n    computed as sigmoid(g_ij) ∈ (0, 1).  The effective (pruned) weight\n    used in the forward pass is:\n\n        pruned_weight_ij = w_ij * sigmoid(g_ij)\n\n    When the gate collapses to ~0 the corresponding weight contributes\n    nothing to the output, effectively removing (\"pruning\") that connection.\n    Gradients flow through both `weight` and `gate_scores` via autograd,\n    so both are updated jointly by the optimizer.\n    \"\"\"\n\n    def __init__(self, in_features: int, out_features: int):\n        super().__init__()\n        self.in_features = in_features\n        self.out_features = out_features\n\n        # Standard weight and bias – same initialisation as nn.Linear\n        self.weight = nn.Parameter(torch.empty(out_features, in_features))\n        self.bias   = nn.Parameter(torch.zeros(out_features))\n\n        # Gate scores: one scalar per weight entry, registered as a parameter\n        # so the optimiser updates them.  Initialised to 0 so that\n        # sigmoid(0) = 0.5 — a neutral starting point.\n        self.gate_scores = nn.Parameter(torch.zeros(out_features, in_features))\n\n        # Weight initialisation (Kaiming uniform, same as nn.Linear default)\n        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))\n        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)\n        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0\n        nn.init.uniform_(self.bias, -bound, bound)\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        # Step 1 – map gate_scores → gates ∈ (0, 1) via sigmoid\n        gates = torch.sigmoid(self.gate_scores)\n\n        # Step 2 – element-wise mask: pruned_weights = weight ⊙ gates\n        pruned_weights = self.weight * gates\n\n        # Step 3 – standard affine transform using the masked weights\n        return F.linear(x, pruned_weights, self.bias)\n\n    def get_gates(self) -> torch.Tensor:\n        \"\"\"Return gate values (detached from graph) for analysis.\"\"\"\n        return torch.sigmoid(self.gate_scores).detach()\n\n    def sparsity(self, threshold: float = 1e-2) -> float:\n        \"\"\"Fraction of gates below `threshold` (i.e., effectively pruned).\"\"\"\n        gates = self.get_gates()\n        return (gates < threshold).float().mean().item()\n\n\n# ─────────────────────────────────────────────────────────────────────────────\n# Network definition\n# ─────────────────────────────────────────────────────────────────────────────\n\nclass SelfPruningNet(nn.Module):\n    \"\"\"\n    A three-hidden-layer feed-forward network built entirely from\n    PrunableLinear layers.\n\n    Input:  CIFAR-10 images flattened to 32×32×3 = 3072 features\n    Output: 10 class logits\n    \"\"\"\n\n    def __init__(self):\n        super().__init__()\n        self.flatten = nn.Flatten()\n\n        self.fc1 = PrunableLinear(3072, 512)\n        self.fc2 = PrunableLinear(512, 256)\n        self.fc3 = PrunableLinear(256, 128)\n        self.fc4 = PrunableLinear(128, 10)\n\n        # Batch-norm layers help stabilise training with aggressive pruning\n        self.bn1 = nn.BatchNorm1d(512)\n        self.bn2 = nn.BatchNorm1d(256)\n        self.bn3 = nn.BatchNorm1d(128)\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        x = self.flatten(x)\n        x = F.relu(self.bn1(self.fc1(x)))\n        x = F.relu(self.bn2(self.fc2(x)))\n        x = F.relu(self.bn3(self.fc3(x)))\n        x = self.fc4(x)\n        return x\n\n    def prunable_layers(self):\n        \"\"\"Iterate over all PrunableLinear layers.\"\"\"\n        for module in self.modules():\n            if isinstance(module, PrunableLinear):\n                yield module\n\n    def sparsity_loss(self) -> torch.Tensor:\n        \"\"\"\n        L1 penalty on all gate values across every PrunableLinear layer.\n\n        Because gates = sigmoid(gate_scores) > 0 always, |gate| = gate, so\n        the L1 norm is simply the sum of all gate values.  This is the term\n        that the optimiser is pushed to minimise, driving gates toward zero.\n        \"\"\"\n        total = torch.tensor(0.0, device=next(self.parameters()).device)\n        count = 0\n        for layer in self.prunable_layers():\n            gates = torch.sigmoid(layer.gate_scores)\n            total = total + gates.sum()\n            count += gates.numel()\n        # Normalise by total number of gates so lambda is scale-invariant\n        return total / max(count, 1)\n\n    def overall_sparsity(self, threshold: float = 1e-2) -> float:\n        \"\"\"\n        Global fraction of gates below `threshold` across the whole network.\n        \"\"\"\n        pruned = total = 0\n        for layer in self.prunable_layers():\n            gates = layer.get_gates()\n            pruned += (gates < threshold).sum().item()\n            total  += gates.numel()\n        return pruned / total if total > 0 else 0.0\n\n    def all_gate_values(self) -> np.ndarray:\n        \"\"\"Concatenate all gate values for histogram plotting.\"\"\"\n        vals = []\n        for layer in self.prunable_layers():\n            vals.append(layer.get_gates().cpu().numpy().ravel())\n        return np.concatenate(vals)\n\n\n# ─────────────────────────────────────────────────────────────────────────────\n# Data loading\n# ─────────────────────────────────────────────────────────────────────────────\n\ndef get_cifar10_loaders(batch_size: int = 128, data_dir: str = \"./data\"):\n    \"\"\"Download CIFAR-10 and return (train_loader, test_loader).\"\"\"\n    mean = (0.4914, 0.4822, 0.4465)\n    std  = (0.2023, 0.1994, 0.2010)\n\n    train_tf = transforms.Compose([\n        transforms.RandomHorizontalFlip(),\n        transforms.RandomCrop(32, padding=4),\n        transforms.ToTensor(),\n        transforms.Normalize(mean, std),\n    ])\n    test_tf = transforms.Compose([\n        transforms.ToTensor(),\n        transforms.Normalize(mean, std),\n    ])\n\n    train_ds = torchvision.datasets.CIFAR10(data_dir, train=True,\n                                             download=True, transform=train_tf)\n    test_ds  = torchvision.datasets.CIFAR10(data_dir, train=False,\n                                             download=True, transform=test_tf)\n\n    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,\n                              num_workers=2, pin_memory=True)\n    test_loader  = DataLoader(test_ds,  batch_size=batch_size, shuffle=False,\n                              num_workers=2, pin_memory=True)\n    return train_loader, test_loader\n\n\n# ─────────────────────────────────────────────────────────────────────────────\n# Part 3: Training and Evaluation\n# ─────────────────────────────────────────────────────────────────────────────\n\ndef train_one_epoch(model, loader, optimizer, device, lam):\n    \"\"\"\n    Run a single training epoch.\n\n    Total Loss = CrossEntropyLoss + λ × SparsityLoss\n    \"\"\"\n    model.train()\n    total_loss = total_cls_loss = total_spar_loss = 0.0\n    correct = total = 0\n\n    for images, labels in loader:\n        images, labels = images.to(device), labels.to(device)\n\n        optimizer.zero_grad()\n        logits = model(images)\n\n        cls_loss  = F.cross_entropy(logits, labels)\n        spar_loss = model.sparsity_loss()\n        loss      = cls_loss + lam * spar_loss\n\n        loss.backward()\n        optimizer.step()\n\n        total_loss      += loss.item()\n        total_cls_loss  += cls_loss.item()\n        total_spar_loss += spar_loss.item()\n\n        preds    = logits.argmax(dim=1)\n        correct += (preds == labels).sum().item()\n        total   += labels.size(0)\n\n    n = len(loader)\n    return {\n        \"loss\":      total_loss      / n,\n        \"cls_loss\":  total_cls_loss  / n,\n        \"spar_loss\": total_spar_loss / n,\n        \"acc\":       correct / total,\n    }\n\n\n@torch.no_grad()\ndef evaluate(model, loader, device):\n    \"\"\"Return test accuracy.\"\"\"\n    model.eval()\n    correct = total = 0\n    for images, labels in loader:\n        images, labels = images.to(device), labels.to(device)\n        preds = model(images).argmax(dim=1)\n        correct += (preds == labels).sum().item()\n        total   += labels.size(0)\n    return correct / total\n\n\ndef run_experiment(lam: float, epochs: int, train_loader, test_loader, device,\n                   out_dir: str = \"outputs\"):\n    \"\"\"\n    Train a full model for the given lambda and return metrics + gate values.\n    \"\"\"\n    print(f\"\\n{'='*60}\")\n    print(f\"  λ = {lam}  |  epochs = {epochs}\")\n    print(f\"{'='*60}\")\n\n    model     = SelfPruningNet().to(device)\n    optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)\n    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n\n    history = []\n    for epoch in range(1, epochs + 1):\n        metrics = train_one_epoch(model, train_loader, optimizer, device, lam)\n        scheduler.step()\n\n        if epoch % 5 == 0 or epoch == epochs:\n            sparsity   = model.overall_sparsity()\n            test_acc   = evaluate(model, test_loader, device)\n            print(f\"  Epoch {epoch:03d} | \"\n                  f\"loss={metrics['loss']:.4f} | \"\n                  f\"train_acc={metrics['acc']:.3f} | \"\n                  f\"test_acc={test_acc:.3f} | \"\n                  f\"sparsity={sparsity:.1%}\")\n            history.append(dict(epoch=epoch, test_acc=test_acc,\n                                sparsity=sparsity, **metrics))\n\n    final_test_acc = evaluate(model, test_loader, device)\n    final_sparsity = model.overall_sparsity()\n    gate_vals      = model.all_gate_values()\n\n    print(f\"\\n  ▶ Final Test Accuracy : {final_test_acc:.4f} \"\n          f\"({final_test_acc*100:.2f}%)\")\n    print(f\"  ▶ Final Sparsity      : {final_sparsity:.4f} \"\n          f\"({final_sparsity*100:.2f}%)\")\n\n    return {\n        \"lam\":           lam,\n        \"test_acc\":      final_test_acc,\n        \"sparsity\":      final_sparsity,\n        \"gate_vals\":     gate_vals,\n        \"history\":       history,\n        \"model\":         model,\n    }\n\n\n# ─────────────────────────────────────────────────────────────────────────────\n# Plotting\n# ─────────────────────────────────────────────────────────────────────────────\n\ndef plot_gate_distribution(results: list, out_path: str = \"gate_distribution.png\"):\n    \"\"\"\n    Plot the gate-value distributions for all lambda values and highlight\n    the best model (highest test accuracy).\n    \"\"\"\n    best = max(results, key=lambda r: r[\"test_acc\"])\n\n    n = len(results)\n    fig, axes = plt.subplots(1, n, figsize=(6 * n, 5), sharey=False)\n    if n == 1:\n        axes = [axes]\n\n    for ax, res in zip(axes, results):\n        gates = res[\"gate_vals\"]\n        lam   = res[\"lam\"]\n\n        # Histogram\n        ax.hist(gates, bins=80, color=\"#4C72B0\", edgecolor=\"none\", alpha=0.85)\n        ax.axvline(1e-2, color=\"red\", linestyle=\"--\", linewidth=1.2,\n                   label=\"Prune threshold (0.01)\")\n        ax.set_xlabel(\"Gate value\", fontsize=12)\n        ax.set_ylabel(\"Count\", fontsize=12)\n        title = f\"λ = {lam}\"\n        if res is best:\n            title += \"  ★ best accuracy\"\n        ax.set_title(title, fontsize=13, fontweight=\"bold\")\n        ax.legend(fontsize=10)\n\n        sparsity_pct = res[\"sparsity\"] * 100\n        acc_pct      = res[\"test_acc\"] * 100\n        ax.text(0.98, 0.97,\n                f\"Sparsity: {sparsity_pct:.1f}%\\nTest Acc: {acc_pct:.1f}%\",\n                transform=ax.transAxes, fontsize=11,\n                verticalalignment=\"top\", horizontalalignment=\"right\",\n                bbox=dict(boxstyle=\"round,pad=0.3\", facecolor=\"lightyellow\",\n                          edgecolor=\"gray\"))\n\n    fig.suptitle(\"Distribution of Gate Values After Training\", fontsize=15,\n                 fontweight=\"bold\", y=1.02)\n    plt.tight_layout()\n    plt.savefig(out_path, dpi=150, bbox_inches=\"tight\")\n    plt.close()\n    print(f\"\\n  [✓] Gate distribution plot saved → {out_path}\")\n\n\ndef plot_training_curves(results: list, out_path: str = \"training_curves.png\"):\n    \"\"\"\n    Loss and sparsity over epochs for each lambda.\n    \"\"\"\n    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\n\n    colours = [\"#1f77b4\", \"#ff7f0e\", \"#2ca02c\", \"#d62728\"]\n    for i, res in enumerate(results):\n        hist   = res[\"history\"]\n        epochs = [h[\"epoch\"] for h in hist]\n        losses = [h[\"loss\"] for h in hist]\n        spars  = [h[\"sparsity\"] * 100 for h in hist]\n        c      = colours[i % len(colours)]\n\n        ax1.plot(epochs, losses, marker=\"o\", color=c, label=f\"λ={res['lam']}\")\n        ax2.plot(epochs, spars,  marker=\"s\", color=c, label=f\"λ={res['lam']}\")\n\n    ax1.set_xlabel(\"Epoch\"); ax1.set_ylabel(\"Total Loss\")\n    ax1.set_title(\"Training Loss vs Epoch\"); ax1.legend()\n\n    ax2.set_xlabel(\"Epoch\"); ax2.set_ylabel(\"Sparsity (%)\")\n    ax2.set_title(\"Sparsity Level vs Epoch\"); ax2.legend()\n\n    plt.tight_layout()\n    plt.savefig(out_path, dpi=150, bbox_inches=\"tight\")\n    plt.close()\n    print(f\"  [✓] Training curves saved → {out_path}\")\n\n\n# ─────────────────────────────────────────────────────────────────────────────\n# Main\n# ─────────────────────────────────────────────────────────────────────────────\n\ndef main():\n    # ── Configuration ────────────────────────────────────────────────────────\n    LAMBDAS    = [1e-4, 1e-3, 5e-3]   # low / medium / high sparsity pressure\n    EPOCHS     = 30                    # increase for production runs\n    BATCH_SIZE = 128\n    OUT_DIR    = \"outputs\"\n    os.makedirs(OUT_DIR, exist_ok=True)\n\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(f\"\\nUsing device: {device}\")\n\n    # ── Data ─────────────────────────────────────────────────────────────────\n    train_loader, test_loader = get_cifar10_loaders(BATCH_SIZE)\n\n    # ── Experiments ──────────────────────────────────────────────────────────\n    all_results = []\n    for lam in LAMBDAS:\n        res = run_experiment(lam, EPOCHS, train_loader, test_loader, device,\n                             out_dir=OUT_DIR)\n        all_results.append(res)\n\n    # ── Summary table ────────────────────────────────────────────────────────\n    print(\"\\n\" + \"─\" * 55)\n    print(f\"{'Lambda':>10} | {'Test Acc (%)':>12} | {'Sparsity (%)':>12}\")\n    print(\"─\" * 55)\n    for r in all_results:\n        print(f\"{r['lam']:>10.4f} | {r['test_acc']*100:>12.2f} | \"\n              f\"{r['sparsity']*100:>12.2f}\")\n    print(\"─\" * 55)\n\n    # ── Save results JSON ────────────────────────────────────────────────────\n    summary = [\n        {\"lambda\": r[\"lam\"],\n         \"test_acc_pct\": round(r[\"test_acc\"] * 100, 2),\n         \"sparsity_pct\": round(r[\"sparsity\"] * 100, 2)}\n        for r in all_results\n    ]\n    json_path = os.path.join(OUT_DIR, \"results_summary.json\")\n    with open(json_path, \"w\") as f:\n        json.dump(summary, f, indent=2)\n    print(f\"\\n  [✓] Results summary saved → {json_path}\")\n\n    # ── Plots ────────────────────────────────────────────────────────────────\n    gate_plot_path   = os.path.join(OUT_DIR, \"gate_distribution.png\")\n    curves_plot_path = os.path.join(OUT_DIR, \"training_curves.png\")\n\n    plot_gate_distribution(all_results, gate_plot_path)\n    plot_training_curves(all_results, curves_plot_path)\n\n    print(\"\\n[✓] All experiments complete.\\n\")\n\n\nif __name__ == \"__main__\":\n    main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-24T05:35:31.610238Z","iopub.execute_input":"2026-04-24T05:35:31.610988Z","iopub.status.idle":"2026-04-24T05:52:45.787029Z","shell.execute_reply.started":"2026-04-24T05:35:31.610949Z","shell.execute_reply":"2026-04-24T05:52:45.786298Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport subprocess, sys\n\n# Check GPU\nprint('PyTorch version :', torch.__version__)\nprint('CUDA available  :', torch.cuda.is_available())\nif torch.cuda.is_available():\n    print('GPU             :', torch.cuda.get_device_name(0))\nelse:\n    print('⚠️  No GPU detected — go to Settings → Accelerator → GPU T4 x2')\n\n# torchvision is pre-installed on Kaggle, just confirm\nimport torchvision\nprint('torchvision     :', torchvision.__version__)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-24T05:56:06.884357Z","iopub.execute_input":"2026-04-24T05:56:06.885073Z","iopub.status.idle":"2026-04-24T05:56:06.891746Z","shell.execute_reply.started":"2026-04-24T05:56:06.885028Z","shell.execute_reply":"2026-04-24T05:56:06.890897Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport math\nimport json\nimport numpy as np\nimport matplotlib\nmatplotlib.use('Agg')\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader\nimport torchvision\nimport torchvision.transforms as transforms\n\nos.makedirs('/kaggle/working/outputs', exist_ok=True)\nprint('✅ All imports successful')\nprint('📁 Output folder ready: /kaggle/working/outputs')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-24T05:56:37.625984Z","iopub.execute_input":"2026-04-24T05:56:37.626699Z","iopub.status.idle":"2026-04-24T05:56:37.632948Z","shell.execute_reply.started":"2026-04-24T05:56:37.626666Z","shell.execute_reply":"2026-04-24T05:56:37.632080Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"\n**Part 1: PrunableLinear LayerCustom linear layer with learnable `gate_scores`. Each weight is multiplied by `sigmoid(gate_score)` s********o the network can learn to zero out unimportant connections.******\n","metadata":{}},{"cell_type":"code","source":"class PrunableLinear(nn.Module):\n    \"\"\"\n    Part 1: Custom linear layer with learnable gate parameters.\n    \n    Forward pass:\n        gates          = sigmoid(gate_scores)          # values in (0, 1)\n        pruned_weights = weight * gates                # element-wise mask\n        output         = pruned_weights @ x + bias     # standard linear op\n    \"\"\"\n    def __init__(self, in_features: int, out_features: int):\n        super().__init__()\n        self.in_features  = in_features\n        self.out_features = out_features\n\n        # ── Standard parameters ───────────────────────────────────────────\n        self.weight = nn.Parameter(torch.empty(out_features, in_features))\n        self.bias   = nn.Parameter(torch.zeros(out_features))\n\n        # ── Gate scores (same shape as weight) ────────────────────────────\n        # Initialised to 0 → sigmoid(0) = 0.5 (neutral start)\n        # Registered as nn.Parameter so optimiser updates them\n        self.gate_scores = nn.Parameter(torch.zeros(out_features, in_features))\n\n        # Kaiming uniform init (same as nn.Linear default)\n        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))\n        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)\n        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0\n        nn.init.uniform_(self.bias, -bound, bound)\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        # Step 1: gate_scores → gates ∈ (0,1) via sigmoid\n        gates = torch.sigmoid(self.gate_scores)\n        # Step 2: element-wise mask  pruned_weights = weight ⊙ gates\n        pruned_weights = self.weight * gates\n        # Step 3: standard linear operation with masked weights\n        return F.linear(x, pruned_weights, self.bias)\n\n    def get_gates(self) -> torch.Tensor:\n        \"\"\"Gate values detached from computation graph (for analysis).\"\"\"\n        return torch.sigmoid(self.gate_scores).detach()\n\n    def sparsity(self, threshold: float = 1e-2) -> float:\n        \"\"\"Fraction of gates below threshold = effectively pruned.\"\"\"\n        return (self.get_gates() < threshold).float().mean().item()\n\n\n# Quick sanity check\nlayer = PrunableLinear(4, 3)\nx_test = torch.randn(2, 4)\nout = layer(x_test)\nprint('✅ PrunableLinear forward pass OK — output shape:', out.shape)\nprint('   gate_scores shape:', layer.gate_scores.shape)\nprint('   gate values range: [{:.3f}, {:.3f}]'.format(\n    layer.get_gates().min().item(), layer.get_gates().max().item()))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-24T05:57:23.004544Z","iopub.execute_input":"2026-04-24T05:57:23.005276Z","iopub.status.idle":"2026-04-24T05:57:23.016080Z","shell.execute_reply.started":"2026-04-24T05:57:23.005241Z","shell.execute_reply":"2026-04-24T05:57:23.015385Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"\n Network + Part 2: Sparsity Loss**Total Loss = CrossEntropyLoss + λ × SparsityLoss**SparsityLoss = mean of all sigmoid(gate_scores) across every layer = L1 norm on gates.The L1 penalty has a constant gradient magnitude, pushing unimportant gates to exactly 0.","metadata":{}},{"cell_type":"code","source":"class SelfPruningNet(nn.Module):\n    \"\"\"\n    4-layer feed-forward network using only PrunableLinear layers.\n    Architecture: 3072 → 512 → 256 → 128 → 10\n    \"\"\"\n    def __init__(self):\n        super().__init__()\n        self.flatten = nn.Flatten()\n\n        # All linear layers are prunable\n        self.fc1 = PrunableLinear(3072, 512)\n        self.fc2 = PrunableLinear(512,  256)\n        self.fc3 = PrunableLinear(256,  128)\n        self.fc4 = PrunableLinear(128,  10)\n\n        # BatchNorm stabilises training under aggressive pruning\n        self.bn1 = nn.BatchNorm1d(512)\n        self.bn2 = nn.BatchNorm1d(256)\n        self.bn3 = nn.BatchNorm1d(128)\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        x = self.flatten(x)                    # 32×32×3 → 3072\n        x = F.relu(self.bn1(self.fc1(x)))      # → 512\n        x = F.relu(self.bn2(self.fc2(x)))      # → 256\n        x = F.relu(self.bn3(self.fc3(x)))      # → 128\n        x = self.fc4(x)                        # → 10 logits\n        return x\n\n    def prunable_layers(self):\n        for m in self.modules():\n            if isinstance(m, PrunableLinear):\n                yield m\n\n    # ── Part 2: Sparsity Loss ─────────────────────────────────────────────\n    def sparsity_loss(self) -> torch.Tensor:\n        \"\"\"\n        Part 2 — L1 penalty on all gate values.\n        \n        Since sigmoid(gate_scores) > 0 always:\n            L1 norm = sum of gate values  (no abs needed)\n        \n        Normalised by total gate count so λ is scale-invariant.\n        \"\"\"\n        total = torch.tensor(0.0, device=next(self.parameters()).device)\n        count = 0\n        for layer in self.prunable_layers():\n            gates  = torch.sigmoid(layer.gate_scores)   # gates ∈ (0,1)\n            total  = total + gates.sum()                # L1 norm\n            count += gates.numel()\n        return total / max(count, 1)                    # normalised\n\n    def overall_sparsity(self, threshold: float = 1e-2) -> float:\n        \"\"\"Global % of gates below threshold across all layers.\"\"\"\n        pruned = total = 0\n        for layer in self.prunable_layers():\n            gates   = layer.get_gates()\n            pruned += (gates < threshold).sum().item()\n            total  += gates.numel()\n        return pruned / total if total > 0 else 0.0\n\n    def all_gate_values(self) -> np.ndarray:\n        \"\"\"All gate values concatenated for histogram plotting.\"\"\"\n        return np.concatenate([\n            layer.get_gates().cpu().numpy().ravel()\n            for layer in self.prunable_layers()\n        ])\n\n\n# Sanity check\nnet = SelfPruningNet()\ndummy = torch.randn(4, 3, 32, 32)\nlogits = net(dummy)\nsl = net.sparsity_loss()\nprint('✅ SelfPruningNet forward pass OK — logits shape:', logits.shape)\nprint('   Initial sparsity loss (should be ~0.5):', round(sl.item(), 4))\nprint('   Total gate parameters:', sum(\n    l.gate_scores.numel() for l in net.prunable_layers()))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-24T05:58:02.112452Z","iopub.execute_input":"2026-04-24T05:58:02.112737Z","iopub.status.idle":"2026-04-24T05:58:02.167264Z","shell.execute_reply.started":"2026-04-24T05:58:02.112714Z","shell.execute_reply":"2026-04-24T05:58:02.166450Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 📂 Cell 5: Load CIFAR-10 Dataset","metadata":{}},{"cell_type":"code","source":"def get_cifar10_loaders(batch_size: int = 128):\n    \"\"\"\n    Load CIFAR-10 from Kaggle's local dataset path.\n    Falls back to auto-download if not found.\n    \"\"\"\n    mean = (0.4914, 0.4822, 0.4465)\n    std  = (0.2023, 0.1994, 0.2010)\n\n    train_tf = transforms.Compose([\n        transforms.RandomHorizontalFlip(),\n        transforms.RandomCrop(32, padding=4),\n        transforms.ToTensor(),\n        transforms.Normalize(mean, std),\n    ])\n    test_tf = transforms.Compose([\n        transforms.ToTensor(),\n        transforms.Normalize(mean, std),\n    ])\n\n    # Use Kaggle's pre-loaded CIFAR-10 if available, else download\n    kaggle_path = '/kaggle/input/cifar-10-python'\n    data_dir = kaggle_path if os.path.exists(kaggle_path) else '/kaggle/working/data'\n    print(f'📂 Using data from: {data_dir}')\n\n    train_ds = torchvision.datasets.CIFAR10(\n        data_dir, train=True, download=True, transform=train_tf)\n    test_ds  = torchvision.datasets.CIFAR10(\n        data_dir, train=False, download=True, transform=test_tf)\n\n    train_loader = DataLoader(train_ds, batch_size=batch_size,\n                              shuffle=True,  num_workers=2, pin_memory=True)\n    test_loader  = DataLoader(test_ds,  batch_size=batch_size,\n                              shuffle=False, num_workers=2, pin_memory=True)\n\n    print(f'✅ Train samples: {len(train_ds):,}  |  Test samples: {len(test_ds):,}')\n    return train_loader, test_loader\n\n\ntrain_loader, test_loader = get_cifar10_loaders(batch_size=128)\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f'🖥️  Training on: {device}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-24T06:00:45.208792Z","iopub.execute_input":"2026-04-24T06:00:45.209614Z","iopub.status.idle":"2026-04-24T06:00:46.830782Z","shell.execute_reply.started":"2026-04-24T06:00:45.209581Z","shell.execute_reply":"2026-04-24T06:00:46.830126Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 🏋️ Cell 6 — Part 3: Training Loop & Evaluation Functions","metadata":{}},{"cell_type":"code","source":"def train_one_epoch(model, loader, optimizer, device, lam):\n    \"\"\"\n    Part 3 — One training epoch.\n    Total Loss = CrossEntropyLoss + λ × SparsityLoss\n    \"\"\"\n    model.train()\n    total_loss = cls_loss_sum = spar_loss_sum = 0.0\n    correct = total = 0\n\n    for images, labels in loader:\n        images, labels = images.to(device), labels.to(device)\n\n        optimizer.zero_grad()\n        logits = model(images)\n\n        # ── Part 2 loss applied here ──────────────────────────────────────\n        cls_loss  = F.cross_entropy(logits, labels)   # classification loss\n        spar_loss = model.sparsity_loss()             # L1 gate penalty\n        loss      = cls_loss + lam * spar_loss        # total loss\n        # ─────────────────────────────────────────────────────────────────\n\n        loss.backward()   # gradients flow through weights AND gate_scores\n        optimizer.step()\n\n        total_loss    += loss.item()\n        cls_loss_sum  += cls_loss.item()\n        spar_loss_sum += spar_loss.item()\n\n        preds    = logits.argmax(dim=1)\n        correct += (preds == labels).sum().item()\n        total   += labels.size(0)\n\n    n = len(loader)\n    return {\n        'loss':      total_loss    / n,\n        'cls_loss':  cls_loss_sum  / n,\n        'spar_loss': spar_loss_sum / n,\n        'acc':       correct / total,\n    }\n\n\n@torch.no_grad()\ndef evaluate(model, loader, device):\n    \"\"\"Return test accuracy.\"\"\"\n    model.eval()\n    correct = total = 0\n    for images, labels in loader:\n        images, labels = images.to(device), labels.to(device)\n        preds   = model(images).argmax(dim=1)\n        correct += (preds == labels).sum().item()\n        total   += labels.size(0)\n    return correct / total\n\n\nprint('✅ Training functions defined')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-24T06:01:18.385308Z","iopub.execute_input":"2026-04-24T06:01:18.385642Z","iopub.status.idle":"2026-04-24T06:01:18.393818Z","shell.execute_reply.started":"2026-04-24T06:01:18.385615Z","shell.execute_reply":"2026-04-24T06:01:18.392990Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 🚀 Cell 7: Run All 3 Lambda Experiments\n","metadata":{}},{"cell_type":"code","source":"# ── Configuration ─────────────────────────────────────────────────────────\nLAMBDAS = [1e-4, 1e-3, 5e-3]   # low / medium / high sparsity pressure\nEPOCHS  = 30                    # increase to 60+ for better accuracy\n\nall_results = []\n\nfor lam in LAMBDAS:\n    print(f\"\\n{'='*60}\")\n    print(f\"  λ = {lam}  |  epochs = {EPOCHS}  |  device = {device}\")\n    print(f\"{'='*60}\")\n\n    model     = SelfPruningNet().to(device)\n    optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)\n    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)\n\n    history = []\n    for epoch in range(1, EPOCHS + 1):\n        metrics = train_one_epoch(model, train_loader, optimizer, device, lam)\n        scheduler.step()\n\n        if epoch % 5 == 0 or epoch == EPOCHS:\n            sparsity = model.overall_sparsity()\n            test_acc = evaluate(model, test_loader, device)\n            print(f\"  Epoch {epoch:03d} | \"\n                  f\"loss={metrics['loss']:.4f} | \"\n                  f\"train_acc={metrics['acc']:.3f} | \"\n                  f\"test_acc={test_acc:.3f} | \"\n                  f\"sparsity={sparsity:.1%}\")\n            history.append(dict(epoch=epoch, test_acc=test_acc,\n                                sparsity=sparsity, **metrics))\n\n    final_test_acc = evaluate(model, test_loader, device)\n    final_sparsity = model.overall_sparsity()\n\n    print(f\"\\n  ▶ Final Test Accuracy : {final_test_acc*100:.2f}%\")\n    print(f\"  ▶ Final Sparsity      : {final_sparsity*100:.2f}%\")\n\n    all_results.append({\n        'lam':       lam,\n        'test_acc':  final_test_acc,\n        'sparsity':  final_sparsity,\n        'gate_vals': model.all_gate_values(),\n        'history':   history,\n    })\n\nprint(\"\\n✅ All experiments complete!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-24T06:11:08.502625Z","iopub.execute_input":"2026-04-24T06:11:08.503285Z","iopub.status.idle":"2026-04-24T06:27:57.272959Z","shell.execute_reply.started":"2026-04-24T06:11:08.503248Z","shell.execute_reply":"2026-04-24T06:27:57.272113Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 📊 Cell 8: Results Summary Table","metadata":{}},{"cell_type":"code","source":"print(\"\\n\" + \"─\"*55)\nprint(f\"{'Lambda':>10} | {'Test Acc (%)':>12} | {'Sparsity (%)':>12}\")\nprint(\"─\"*55)\nfor r in all_results:\n    print(f\"{r['lam']:>10.4f} | {r['test_acc']*100:>12.2f} | {r['sparsity']*100:>12.2f}\")\nprint(\"─\"*55)\n\n# Save JSON\nsummary = [\n    {'lambda': r['lam'],\n     'test_acc_pct': round(r['test_acc'] * 100, 2),\n     'sparsity_pct': round(r['sparsity'] * 100, 2)}\n    for r in all_results\n]\nwith open('/kaggle/working/outputs/results_summary.json', 'w') as f:\n    json.dump(summary, f, indent=2)\nprint(\"\\n✅ Saved → /kaggle/working/outputs/results_summary.json\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-24T06:27:57.645378Z","iopub.execute_input":"2026-04-24T06:27:57.646035Z","iopub.status.idle":"2026-04-24T06:27:57.652726Z","shell.execute_reply.started":"2026-04-24T06:27:57.646003Z","shell.execute_reply":"2026-04-24T06:27:57.651974Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 📈 Cell 9: Gate Distribution Plot\nA successful result shows a **big spike at 0** (pruned weights) and a **separate cluster > 0** (kept weights).","metadata":{}},{"cell_type":"code","source":"best = max(all_results, key=lambda r: r['test_acc'])\nn    = len(all_results)\n\nfig, axes = plt.subplots(1, n, figsize=(6*n, 5))\nif n == 1:\n    axes = [axes]\n\nfor ax, res in zip(axes, all_results):\n    ax.hist(res['gate_vals'], bins=80, color='#4C72B0',\n            edgecolor='none', alpha=0.85)\n    ax.axvline(1e-2, color='red', linestyle='--',\n               linewidth=1.5, label='Prune threshold (0.01)')\n    ax.set_xlabel('Gate value', fontsize=12)\n    ax.set_ylabel('Count',      fontsize=12)\n    title = f\"λ = {res['lam']}\"\n    if res is best:\n        title += '  ★ best accuracy'\n    ax.set_title(title, fontsize=13, fontweight='bold')\n    ax.legend(fontsize=10)\n    ax.text(0.98, 0.97,\n            f\"Sparsity: {res['sparsity']*100:.1f}%\\nTest Acc: {res['test_acc']*100:.1f}%\",\n            transform=ax.transAxes, fontsize=11,\n            va='top', ha='right',\n            bbox=dict(boxstyle='round,pad=0.3',\n                      facecolor='lightyellow', edgecolor='gray'))\n\nfig.suptitle('Distribution of Gate Values After Training',\n             fontsize=15, fontweight='bold', y=1.02)\nplt.tight_layout()\nout_path = '/kaggle/working/outputs/gate_distribution.png'\nplt.savefig(out_path, dpi=150, bbox_inches='tight')\nplt.show()\nprint(f'✅ Saved → {out_path}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-24T06:28:13.356354Z","iopub.execute_input":"2026-04-24T06:28:13.356770Z","iopub.status.idle":"2026-04-24T06:28:14.399862Z","shell.execute_reply.started":"2026-04-24T06:28:13.356741Z","shell.execute_reply":"2026-04-24T06:28:14.399023Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 📉 Cell 10: Training Loss & Sparsity Curves","metadata":{}},{"cell_type":"code","source":"fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))\ncolours = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728']\n\nfor i, res in enumerate(all_results):\n    hist   = res['history']\n    epochs = [h['epoch']        for h in hist]\n    losses = [h['loss']         for h in hist]\n    spars  = [h['sparsity']*100 for h in hist]\n    c = colours[i % len(colours)]\n\n    ax1.plot(epochs, losses, marker='o', color=c, label=f\"λ={res['lam']}\")\n    ax2.plot(epochs, spars,  marker='s', color=c, label=f\"λ={res['lam']}\")\n\nax1.set_xlabel('Epoch');        ax1.set_ylabel('Total Loss')\nax1.set_title('Training Loss vs Epoch'); ax1.legend()\n\nax2.set_xlabel('Epoch');        ax2.set_ylabel('Sparsity (%)')\nax2.set_title('Sparsity Level vs Epoch'); ax2.legend()\n\nplt.tight_layout()\nout_path = '/kaggle/working/outputs/training_curves.png'\nplt.savefig(out_path, dpi=150, bbox_inches='tight')\nplt.show()\nprint(f'✅ Saved → {out_path}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-24T06:28:39.601151Z","iopub.execute_input":"2026-04-24T06:28:39.601889Z","iopub.status.idle":"2026-04-24T06:28:40.051718Z","shell.execute_reply.started":"2026-04-24T06:28:39.601856Z","shell.execute_reply":"2026-04-24T06:28:40.051072Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## ✅ Cell 11: Final Output Files","metadata":{}},{"cell_type":"code","source":"print(\"\\n📁 Output files generated:\")\nfor f in sorted(os.listdir('/kaggle/working/outputs')):\n    path = f'/kaggle/working/outputs/{f}'\n    size = os.path.getsize(path)\n    print(f\"   {f:<35} {size/1024:.1f} KB\")\n\nprint(\"\\n🎉 Done! View outputs in the right panel under Output → /kaggle/working/outputs/\")\nprint(\"   Save Version (top-right) to keep your results permanently.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-24T06:29:12.437302Z","iopub.execute_input":"2026-04-24T06:29:12.437580Z","iopub.status.idle":"2026-04-24T06:29:12.443388Z","shell.execute_reply.started":"2026-04-24T06:29:12.437556Z","shell.execute_reply":"2026-04-24T06:29:12.442488Z"}},"outputs":[],"execution_count":null}]}