{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  },
  "colab": {
   "provenance": []
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-0000",
   "metadata": {},
   "source": [
    "# Minibatches, Datasets & DataLoaders\n",
    "\n",
    "Why we train on small batches, how Dataset and DataLoader feed them to the model, and the full modern training recipe with validation, checkpointing, and early stopping.\n",
    "\n",
    "*Part of the free [Deep Learning with PyTorch](https://ramadnsyh.dev/courses/deep-learning) course by [Muhammad Ramadiansyah](https://ramadnsyh.dev). This notebook is generated from the interactive lesson — [read it online](https://ramadnsyh.dev/courses/deep-learning/minibatches-dataloaders).*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0001",
   "metadata": {},
   "source": [
    "So far we've pushed the whole training set through the network in one go.\n",
    "That works for a thousand points; it collapses for a million images — they\n",
    "simply don't fit in memory. The fix is to train on small **minibatches**, and\n",
    "PyTorch has a dedicated data pipeline for it: `Dataset` and `DataLoader`. By\n",
    "the end of this lesson you'll have a complete, reusable training recipe —\n",
    "validation per epoch, best-model checkpointing, and early stopping — that\n",
    "carries you through the rest of the course."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0002",
   "metadata": {},
   "source": [
    "## Why minibatches?\n",
    "\n",
    "Instead of computing the loss on all N examples per update, split the data\n",
    "into chunks of, say, 64 and update after each chunk. Three wins:\n",
    "\n",
    "1. **Memory.** Only one batch lives on the GPU at a time. A dataset of any\n",
    "   size trains on a fixed memory budget.\n",
    "2. **Faster convergence.** With batches of 64 and 64,000 examples, you get\n",
    "   1,000 weight updates per epoch instead of one. Each update uses a noisier\n",
    "   gradient estimate, but a thousand decent steps beat one perfect step.\n",
    "3. **Noise as a feature.** Batch gradients only *approximate* the full\n",
    "   gradient, so the path downhill jitters. That jitter acts as a mild\n",
    "   regularizer and helps the optimizer escape sharp, poorly-generalizing\n",
    "   minima.\n",
    "\n",
    "Three words you'll now use precisely:\n",
    "\n",
    "- **Batch** — one chunk of examples (its size is the *batch size*).\n",
    "- **Iteration** — one forward-backward-step cycle on one batch.\n",
    "- **Epoch** — one full pass over the dataset. With 1,000 samples and batch\n",
    "  size 64, an epoch is ceil(1000 / 64) = 16 iterations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0003",
   "metadata": {},
   "source": [
    "## Dataset and DataLoader\n",
    "\n",
    "PyTorch splits the job in two. A **`Dataset`** answers two questions: *how\n",
    "many examples are there?* (`__len__`) and *give me example i* (`__getitem__`).\n",
    "A **`DataLoader`** wraps a dataset and handles batching, shuffling, and\n",
    "parallel loading. Writing a custom dataset is just the OOP you learned last\n",
    "lesson:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0004",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader, TensorDataset\n",
    "\n",
    "class MyDataset(Dataset):\n",
    "    def __init__(self, X, y):\n",
    "        self.X = torch.tensor(X, dtype=torch.float32)\n",
    "        self.y = torch.tensor(y, dtype=torch.long)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.X)              # how many samples?\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        return self.X[i], self.y[i]     # one (features, label) pair\n",
    "\n",
    "# For tensors that already exist, TensorDataset does the same thing:\n",
    "import numpy as np\n",
    "rng = np.random.default_rng(42)\n",
    "X, y = rng.normal(size=(1000, 8)), rng.integers(0, 2, 1000)\n",
    "\n",
    "train_set = TensorDataset(\n",
    "    torch.tensor(X, dtype=torch.float32),\n",
    "    torch.tensor(y, dtype=torch.long),\n",
    ")\n",
    "\n",
    "trainloader = DataLoader(train_set, batch_size=64, shuffle=True)\n",
    "\n",
    "for xb, yb in trainloader:              # 16 iterations per epoch\n",
    "    print(xb.shape, yb.shape)           # torch.Size([64, 8]) torch.Size([64])\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0005",
   "metadata": {},
   "source": [
    "Custom `Dataset` classes earn their keep when `__getitem__` does real work —\n",
    "loading an image file from disk, applying transforms — so that only the\n",
    "current batch is ever in memory.\n",
    "\n",
    "Two `DataLoader` knobs matter most:\n",
    "\n",
    "- **`shuffle=True`** for training, always. If the data is sorted (all class 0,\n",
    "  then all class 1...), each batch is lopsided and the gradients lurch from\n",
    "  one class's preferences to the other's. Shuffling each epoch keeps batches\n",
    "  representative. For validation, `shuffle=False` — order doesn't affect a\n",
    "  metric, and reproducibility is nice.\n",
    "- **`batch_size`** trades speed for noise. Small batches (8–32): noisy\n",
    "  gradients, more regularization, slower wall-clock. Large batches (256+):\n",
    "  smooth gradients, better hardware utilization, but more memory and\n",
    "  sometimes worse generalization. **64 or 128 is a sensible default**; cut it\n",
    "  if you hit out-of-memory errors."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0006",
   "metadata": {},
   "source": [
    "## The full modern recipe\n",
    "\n",
    "Minibatches change the loop's shape: an inner loop over batches, nested in an\n",
    "outer loop over epochs — and after each epoch, an evaluation pass over a\n",
    "**validation set** the model never trains on. Comparing the two curves is how\n",
    "you diagnose overfitting, and the validation score drives two more upgrades:\n",
    "\n",
    "- **Checkpointing** — whenever validation improves, save the weights with\n",
    "  `torch.save(model.state_dict(), path)`. Training can wander into\n",
    "  overfitting; the best model is safely on disk.\n",
    "- **Early stopping** — if validation hasn't improved for `patience` epochs in\n",
    "  a row, stop. No babysitting, no wasted compute, no guessing the right\n",
    "  number of epochs in advance.\n",
    "\n",
    "Here's the whole thing as a reusable function — copy it into your projects:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0007",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "def run_epoch(model, loader, criterion, optimizer=None):\n",
    "    \"\"\"One pass over loader. Trains if optimizer is given, else evaluates.\"\"\"\n",
    "    training = optimizer is not None\n",
    "    model.train() if training else model.eval()\n",
    "\n",
    "    total_loss, correct, n = 0.0, 0, 0\n",
    "    with torch.enable_grad() if training else torch.inference_mode():\n",
    "        for xb, yb in loader:\n",
    "            xb, yb = xb.to(device), yb.to(device)\n",
    "            output = model(xb)\n",
    "            loss = criterion(output, yb)\n",
    "\n",
    "            if training:\n",
    "                optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "\n",
    "            total_loss += loss.item() * len(xb)      # weight by batch size\n",
    "            correct += (output.argmax(1) == yb).sum().item()\n",
    "            n += len(xb)\n",
    "    return total_loss / n, correct / n\n",
    "\n",
    "def train(model, trainloader, valloader, criterion, optimizer,\n",
    "          max_epochs=200, patience=10, path=\"best_model.pth\"):\n",
    "    best_val_loss, wait = float(\"inf\"), 0\n",
    "\n",
    "    for epoch in range(1, max_epochs + 1):\n",
    "        train_loss, train_acc = run_epoch(model, trainloader, criterion, optimizer)\n",
    "        val_loss, val_acc = run_epoch(model, valloader, criterion)\n",
    "\n",
    "        if val_loss < best_val_loss:                 # improvement: checkpoint\n",
    "            best_val_loss, wait = val_loss, 0\n",
    "            torch.save(model.state_dict(), path)\n",
    "        else:                                        # no improvement\n",
    "            wait += 1\n",
    "            if wait >= patience:\n",
    "                print(f\"early stopping at epoch {epoch}\")\n",
    "                break\n",
    "\n",
    "        print(f\"epoch {epoch:3d} | train loss {train_loss:.4f} acc {train_acc:.3f}\"\n",
    "              f\" | val loss {val_loss:.4f} acc {val_acc:.3f}\")\n",
    "\n",
    "    model.load_state_dict(torch.load(path))          # restore the best weights\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0008",
   "metadata": {},
   "source": [
    "Details worth noticing:\n",
    "\n",
    "- The last batch is usually smaller than `batch_size`, so we accumulate\n",
    "  `loss.item() * len(xb)` and divide by the total count — a plain average of\n",
    "  batch losses would weight the final stragglers too heavily.\n",
    "- `run_epoch` serves both phases: pass an optimizer to train, omit it to\n",
    "  evaluate under `inference_mode`.\n",
    "- We monitor **validation loss** and reload the checkpoint at the end, so the\n",
    "  function returns the *best* model, not the last one.\n",
    "\n",
    "And using it is three lines:"
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0009",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "model = nn.Sequential(\n",
    "    nn.Linear(8, 32), nn.ReLU(), nn.Dropout(0.2),\n",
    "    nn.Linear(32, 2),\n",
    ").to(device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "model = train(model, trainloader, valloader, criterion, optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0010",
   "metadata": {},
   "source": [
    "> **Save the state_dict, not the model**\n",
    "> \n",
    "> `torch.save(model, path)` pickles the whole Python object and breaks the\n",
    "> moment your class definition or PyTorch version changes. The robust pattern is\n",
    "> `torch.save(model.state_dict(), path)`, then later rebuild the architecture in\n",
    "> code and `model.load_state_dict(torch.load(path))`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0011",
   "metadata": {},
   "source": [
    "### 🏋️ Exercise — Put the recipe to work\n",
    "\n",
    "val loss {val_loss:.4f} | val acc {val_acc:.3f}\")\n",
    "# typically ~0.97-0.99 accuracy, stopping well before 200 epochs\n",
    "`}\n",
    ">\n",
    "Train a classifier on scikit-learn's `load_breast_cancer` dataset (30\n",
    "features, 2 classes) using the full recipe: scale the features, build\n",
    "train/validation `DataLoader`s (batch size 32, shuffled training only), define\n",
    "a small MLP with dropout, and run the `train()` function with `patience=10`.\n",
    "Report the best validation accuracy and note which epoch early stopping fired\n",
    "at."
   ]
  },
  {
   "cell_type": "code",
   "id": "cell-0012",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your solution here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0013",
   "metadata": {},
   "source": [
    "<details><summary>✅ Show solution</summary>\n",
    "\n",
    "```python\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "from sklearn.datasets import load_breast_cancer\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "torch.manual_seed(42)\n",
    "\n",
    "X, y = load_breast_cancer(return_X_y=True)\n",
    "X_train, X_val, y_train, y_val = train_test_split(\n",
    "    X, y, test_size=0.2, stratify=y, random_state=42)\n",
    "\n",
    "scaler = StandardScaler()\n",
    "X_train = scaler.fit_transform(X_train)\n",
    "X_val = scaler.transform(X_val)\n",
    "\n",
    "def to_dataset(X, y):\n",
    "    return TensorDataset(torch.tensor(X, dtype=torch.float32),\n",
    "                         torch.tensor(y, dtype=torch.long))\n",
    "\n",
    "trainloader = DataLoader(to_dataset(X_train, y_train), batch_size=32, shuffle=True)\n",
    "valloader = DataLoader(to_dataset(X_val, y_val), batch_size=64, shuffle=False)\n",
    "\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(30, 32), nn.ReLU(), nn.Dropout(0.2),\n",
    "    nn.Linear(32, 2),\n",
    ").to(device)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# paste run_epoch and train from the lesson here, then:\n",
    "model = train(model, trainloader, valloader, criterion, optimizer,\n",
    "              max_epochs=200, patience=10)\n",
    "\n",
    "val_loss, val_acc = run_epoch(model, valloader, criterion)\n",
    "print(f\"best model -> val loss {val_loss:.4f} | val acc {val_acc:.3f}\")\n",
    "# typically ~0.97-0.99 accuracy, stopping well before 200 epochs\n",
    "```\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-0014",
   "metadata": {},
   "source": [
    "Next up: a new module — convolutional neural networks, where we finally stop\n",
    "flattening images and let the network see their 2-D structure."
   ]
  }
 ]
}