{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "This is a companion notebook for the book [Deep Learning with Python, Third Edition](https://www.manning.com/books/deep-learning-with-python-third-edition). For readability, it only contains runnable code blocks and section titles, and omits everything else in the book: text paragraphs, figures, and pseudocode.\n\n**If you want to be able to follow what's going on, I recommend reading the notebook side by side with your copy of the book.**\n\nThe book's contents are available online at [deeplearningwithpython.io](https://deeplearningwithpython.io)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "!pip install keras keras-hub --upgrade -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"KERAS_BACKEND\"] = \"jax\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "cellView": "form",
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "# @title\n",
    "import os\n",
    "from IPython.core.magic import register_cell_magic\n",
    "\n",
    "@register_cell_magic\n",
    "def backend(line, cell):\n",
    "    current, required = os.environ.get(\"KERAS_BACKEND\", \"\"), line.split()[-1]\n",
    "    if current == required:\n",
    "        get_ipython().run_cell(cell)\n",
    "    else:\n",
    "        print(\n",
    "            f\"This cell requires the {required} backend. To run it, change KERAS_BACKEND to \"\n",
    "            f\"\\\"{required}\\\" at the top of the notebook, restart the runtime, and rerun the notebook.\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Best practices for the real world"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Getting the most out of your models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Hyperparameter optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Using KerasTuner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "!pip install keras-tuner -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import keras\n",
    "from keras import layers\n",
    "\n",
    "def build_model(hp):\n",
    "    units = hp.Int(name=\"units\", min_value=16, max_value=64, step=16)\n",
    "    model = keras.Sequential(\n",
    "        [\n",
    "            layers.Dense(units, activation=\"relu\"),\n",
    "            layers.Dense(10, activation=\"softmax\"),\n",
    "        ]\n",
    "    )\n",
    "    optimizer = hp.Choice(name=\"optimizer\", values=[\"rmsprop\", \"adam\"])\n",
    "    model.compile(\n",
    "        optimizer=optimizer,\n",
    "        loss=\"sparse_categorical_crossentropy\",\n",
    "        metrics=[\"accuracy\"],\n",
    "    )\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import keras_tuner as kt\n",
    "\n",
    "class SimpleMLP(kt.HyperModel):\n",
    "    def __init__(self, num_classes):\n",
    "        self.num_classes = num_classes\n",
    "\n",
    "    def build(self, hp):\n",
    "        units = hp.Int(name=\"units\", min_value=16, max_value=64, step=16)\n",
    "        model = keras.Sequential(\n",
    "            [\n",
    "                layers.Dense(units, activation=\"relu\"),\n",
    "                layers.Dense(self.num_classes, activation=\"softmax\"),\n",
    "            ]\n",
    "        )\n",
    "        optimizer = hp.Choice(name=\"optimizer\", values=[\"rmsprop\", \"adam\"])\n",
    "        model.compile(\n",
    "            optimizer=optimizer,\n",
    "            loss=\"sparse_categorical_crossentropy\",\n",
    "            metrics=[\"accuracy\"],\n",
    "        )\n",
    "        return model\n",
    "\n",
    "hypermodel = SimpleMLP(num_classes=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "tuner = kt.BayesianOptimization(\n",
    "    build_model,\n",
    "    objective=\"val_accuracy\",\n",
    "    max_trials=20,\n",
    "    executions_per_trial=2,\n",
    "    directory=\"mnist_kt_test\",\n",
    "    overwrite=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "tuner.search_space_summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()\n",
    "x_train = x_train.reshape((-1, 28 * 28)).astype(\"float32\") / 255\n",
    "x_test = x_test.reshape((-1, 28 * 28)).astype(\"float32\") / 255\n",
    "x_train_full = x_train[:]\n",
    "y_train_full = y_train[:]\n",
    "num_val_samples = 10000\n",
    "x_train, x_val = x_train[:-num_val_samples], x_train[-num_val_samples:]\n",
    "y_train, y_val = y_train[:-num_val_samples], y_train[-num_val_samples:]\n",
    "callbacks = [\n",
    "    keras.callbacks.EarlyStopping(monitor=\"val_loss\", patience=5),\n",
    "]\n",
    "tuner.search(\n",
    "    x_train,\n",
    "    y_train,\n",
    "    batch_size=128,\n",
    "    epochs=100,\n",
    "    validation_data=(x_val, y_val),\n",
    "    callbacks=callbacks,\n",
    "    verbose=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "top_n = 4\n",
    "best_hps = tuner.get_best_hyperparameters(top_n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "def get_best_epoch(hp):\n",
    "    model = build_model(hp)\n",
    "    callbacks = [\n",
    "        keras.callbacks.EarlyStopping(\n",
    "            monitor=\"val_loss\", mode=\"min\", patience=10\n",
    "        )\n",
    "    ]\n",
    "    history = model.fit(\n",
    "        x_train,\n",
    "        y_train,\n",
    "        validation_data=(x_val, y_val),\n",
    "        epochs=100,\n",
    "        batch_size=128,\n",
    "        callbacks=callbacks,\n",
    "    )\n",
    "    val_loss_per_epoch = history.history[\"val_loss\"]\n",
    "    best_epoch = val_loss_per_epoch.index(min(val_loss_per_epoch)) + 1\n",
    "    print(f\"Best epoch: {best_epoch}\")\n",
    "    return best_epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "def get_best_trained_model(hp):\n",
    "    best_epoch = get_best_epoch(hp)\n",
    "    model = build_model(hp)\n",
    "    model.fit(\n",
    "        x_train_full, y_train_full, batch_size=128, epochs=int(best_epoch * 1.2)\n",
    "    )\n",
    "    return model\n",
    "\n",
    "best_models = []\n",
    "for hp in best_hps:\n",
    "    model = get_best_trained_model(hp)\n",
    "    model.evaluate(x_test, y_test)\n",
    "    best_models.append(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "best_models = tuner.get_best_models(top_n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### The art of crafting the right search space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### The future of hyperparameter tuning: automated machine learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Model ensembling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Scaling up model training with multiple devices"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Multi-GPU training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Data parallelism: Replicating your model on each GPU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Model parallelism: Splitting your model across multiple GPUs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Distributed training in practice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Getting your hands on two or more GPUs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Using data parallelism with JAX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Using model parallelism with JAX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "###### The DeviceMesh API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "###### The LayoutMap API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### TPU training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Using step fusing to improve TPU utilization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Speeding up training and inference with lower-precision computation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Understanding floating-point precision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Float16 inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Mixed-precision training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Using loss scaling with mixed precision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "##### Beyond mixed precision: float8 training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "#### Faster inference with quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "from keras import ops\n",
    "\n",
    "x = ops.array([[0.1, 0.9], [1.2, -0.8]])\n",
    "kernel = ops.array([[-0.1, -2.2], [1.1, 0.7]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "def abs_max_quantize(value):\n",
    "    abs_max = ops.max(ops.abs(value), keepdims=True)\n",
    "    scale = ops.divide(127, abs_max + 1e-7)\n",
    "    scaled_value = value * scale\n",
    "    scaled_value = ops.clip(ops.round(scaled_value), -127, 127)\n",
    "    scaled_value = ops.cast(scaled_value, dtype=\"int8\")\n",
    "    return scaled_value, scale\n",
    "\n",
    "int_x, x_scale = abs_max_quantize(x)\n",
    "int_kernel, kernel_scale = abs_max_quantize(kernel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "int_y = ops.matmul(int_x, int_kernel)\n",
    "y = ops.cast(int_y, dtype=\"float32\") / (x_scale * kernel_scale)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "ops.matmul(x, kernel)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "chapter18_best-practices-for-the-real-world",
   "private_outputs": false,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}