{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3d5eca2f",
   "metadata": {},
   "source": [
    "# Improving deep networks in Keras 3 and TensorFlow\n",
    "\n",
    "The same ideas the TensorFlow way: Keras 3 layers, `compile`/`fit`, callbacks, and `tf.GradientTape` when you want your own training loop. TensorFlow 2.21 is stable and mostly receives fixes now; Google recommends Keras 3, JAX or PyTorch for new generative-AI work, and Keras 3 code like this also runs on JAX or PyTorch by changing one line.\n",
    "\n",
    "[Open in Google Colab](https://colab.research.google.com/github/mahmoudyosrimahmoud13/inpace-library/blob/main/courses/dl-2-improving-networks/files/dl2_improving_networks_keras.ipynb) · part of the inPace course library. Slides © DeepLearning.AI (CC BY-SA 2.0); this notebook: inPace (CC BY-SA 4.0).\n",
    "\n",
    "**Run it top to bottom** (Runtime → Run all in Colab). Each part says what to look for."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33eb6cb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Works in Google Colab as it is (TensorFlow and Keras are already installed there).\n",
    "# On your own computer:  pip install tensorflow keras scikit-learn scipy matplotlib\n",
    "# One framework per notebook: PyTorch and TensorFlow can clash when loaded into the same Python process.\n",
    "import os\n",
    "os.environ[\"KERAS_BACKEND\"] = \"tensorflow\"   # Keras 3 also runs on \"jax\" or \"torch\": change it here\n",
    "os.environ.setdefault(\"TF_CPP_MIN_LOG_LEVEL\", \"3\")\n",
    "import numpy as np, matplotlib.pyplot as plt\n",
    "import tensorflow as tf, keras\n",
    "keras.utils.set_random_seed(0)\n",
    "print(\"TensorFlow\", tf.__version__, \"| Keras\", keras.__version__, \"| backend:\", keras.backend.backend(),\n",
    "      \"| GPU:\", bool(tf.config.list_physical_devices(\"GPU\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d2883c2",
   "metadata": {},
   "source": [
    "## 0. Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ddf23db",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.datasets import load_digits\n",
    "from sklearn.model_selection import train_test_split\n",
    "d = load_digits()\n",
    "X, Xt_, y, yt_ = train_test_split((d.data / 16.0).astype(\"float32\"), d.target, test_size=0.4, random_state=0)\n",
    "Xd, Xte, yd, yte = train_test_split(Xt_, yt_, test_size=0.5, random_state=0)\n",
    "LOSS = keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "def mlp(width=256, depth=3, dropout=0.0, l2=0.0, bn=False):\n",
    "    layers = [keras.Input((64,))]\n",
    "    for _ in range(depth):\n",
    "        layers.append(keras.layers.Dense(width, kernel_regularizer=keras.regularizers.L2(l2) if l2 else None))\n",
    "        if bn: layers.append(keras.layers.BatchNormalization())\n",
    "        layers += [keras.layers.ReLU(), keras.layers.Dropout(dropout)]\n",
    "    return keras.Sequential(layers + [keras.layers.Dense(10)])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8516cab",
   "metadata": {},
   "source": [
    "## 1. Regularization on 150 examples: L2 and dropout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e143c193",
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, kw in ((\"nothing\", {}), (\"L2 1e-3\", {\"l2\": 1e-3}), (\"dropout 0.5\", {\"dropout\": 0.5})):\n",
    "    keras.utils.set_random_seed(0); m = mlp(512, **kw)\n",
    "    m.compile(keras.optimizers.Adam(1e-3), LOSS, metrics=[\"accuracy\"])\n",
    "    h = m.fit(X[:150], y[:150], validation_data=(Xd, yd), epochs=150, batch_size=50, verbose=0)\n",
    "    plt.plot(h.history[\"val_loss\"], label=f\"{name}: dev acc {h.history['val_accuracy'][-1]:.3f}\")\n",
    "plt.legend(); plt.ylim(0, 1.5); plt.title(\"dev loss\"); plt.xlabel(\"epoch\"); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3a888bc",
   "metadata": {},
   "source": [
    "## 2. Optimizers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a08e33d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "opts = {\"SGD\": lambda: keras.optimizers.SGD(0.05), \"momentum\": lambda: keras.optimizers.SGD(0.05, momentum=0.9),\n",
    "        \"RMSprop\": lambda: keras.optimizers.RMSprop(1e-3), \"Adam\": lambda: keras.optimizers.Adam(1e-3),\n",
    "        \"AdamW\": lambda: keras.optimizers.AdamW(1e-3, weight_decay=0.01)}\n",
    "for name, make in opts.items():\n",
    "    keras.utils.set_random_seed(0); m = mlp(); m.compile(make(), LOSS, metrics=[\"accuracy\"])\n",
    "    h = m.fit(X, y, epochs=25, batch_size=64, verbose=0)\n",
    "    plt.plot(h.history[\"loss\"], label=f\"{name}: test acc {m.evaluate(Xte, yte, verbose=0)[1]:.3f}\")\n",
    "plt.yscale(\"log\"); plt.legend(); plt.title(\"training loss\"); plt.xlabel(\"epoch\"); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0343fcf2",
   "metadata": {},
   "source": [
    "## 3. A schedule, batch norm and early stopping together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f7558a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "steps = 60 * int(np.ceil(len(X) / 64))\n",
    "keras.utils.set_random_seed(0); m = mlp(256, 3, dropout=0.3, bn=True)\n",
    "m.compile(keras.optimizers.AdamW(keras.optimizers.schedules.CosineDecay(1e-4, steps, warmup_target=3e-3, warmup_steps=steps // 20),\n",
    "                                  weight_decay=1e-2), LOSS, metrics=[\"accuracy\"])\n",
    "h = m.fit(X, y, validation_data=(Xd, yd), epochs=60, batch_size=64, verbose=0,\n",
    "          callbacks=[keras.callbacks.EarlyStopping(patience=8, restore_best_weights=True)])\n",
    "print(f\"stopped after {len(h.history['loss'])} epochs | test accuracy {m.evaluate(Xte, yte, verbose=0)[1]:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "004d3c72",
   "metadata": {},
   "source": [
    "## 4. tf.data: the input pipeline\n",
    "Shuffle, batch and prefetch so the GPU never waits for data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4f89c7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = tf.data.Dataset.from_tensor_slices((X, y)).shuffle(1000).batch(64).prefetch(tf.data.AUTOTUNE)\n",
    "xb, yb = next(iter(ds)); print(\"one batch:\", xb.shape, yb.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61afdf2a",
   "metadata": {},
   "source": [
    "## 5. Mixed precision on a GPU\n",
    "In Colab with a GPU: `keras.mixed_precision.set_global_policy(\"mixed_float16\")` before building the model. Keep the last layer in float32."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "022ae1d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "if tf.config.list_physical_devices(\"GPU\"):\n",
    "    keras.mixed_precision.set_global_policy(\"mixed_float16\")\n",
    "print(\"policy:\", keras.mixed_precision.global_policy().name)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}