keras.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "13061b0b",
   "metadata": {},
   "source": [
    "# Simple Neural Network Regression with Keras and JAX"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2526d005",
   "metadata": {},
   "source": [
    "# GPU\n",
    "\n",
    "This lecture was built using a machine with JAX installed and access to a GPU.\n",
    "\n",
    "To run this lecture on [Google Colab](https://colab.research.google.com/), click on the “play” icon top right, select Colab, and set the runtime environment to include a GPU.\n",
    "\n",
    "To run this lecture on your own machine, you need to install [Google JAX](https://github.com/google/jax).\n",
    "\n",
    "In this lecture we show how to implement one-dimensional nonlinear regression\n",
    "using a neural network.\n",
    "\n",
    "We will use the popular deep learning library [Keras](https://keras.io/), which\n",
    "provides a simple and elegant interface to deep learning.\n",
    "\n",
    "The emphasis in Keras is on providing an intuitive API, while the heavy lifting is\n",
    "done by another library.\n",
    "\n",
    "Currently the backend library can be Tensorflow, PyTorch, or JAX.\n",
    "\n",
    "In this lecture we will use JAX.\n",
    "\n",
    "The objective of this lecture is to provide a very simple introduction to deep\n",
    "learning in a regression setting.\n",
    "\n",
    "We begin with some standard imports."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9154be3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c85bcd1",
   "metadata": {},
   "source": [
    "Let’s install Keras."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79e3d2bd",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install keras"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d0a79bb",
   "metadata": {},
   "source": [
    "Now we specify that the desired backend is JAX."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a8a0b41",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['KERAS_BACKEND'] = 'jax'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "319ed973",
   "metadata": {},
   "source": [
    "Next we import some tools from Keras."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6f04e3e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import keras\n",
    "from keras.models import Sequential\n",
    "from keras.layers import Dense"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85fcd00f",
   "metadata": {},
   "source": [
    "## Data\n",
    "\n",
    "First let’s write a function to generate some data.\n",
    "\n",
    "The data has the form\n",
    "\n",
    "$$\n",
    "y_i = f(x_i) + \\epsilon_i,\n",
    "    \\qquad i=1, \\ldots, n\n",
    "$$\n",
    "\n",
    "The map $ f $ is specified inside the function and $ \\epsilon_i $ is an independent\n",
    "draw from a fixed normal distribution.\n",
    "\n",
    "Here’s the function that creates vectors `x` and `y` according to the rule\n",
    "above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe26590a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def generate_data(x_min=0, x_max=5, data_size=400):\n",
    "    x = np.linspace(x_min, x_max, num=data_size)\n",
    "    x = x.reshape(data_size, 1)\n",
    "    ϵ = 0.2 * np.random.randn(*x.shape)\n",
    "    y = x**0.5 + np.sin(x) + ϵ\n",
    "    x, y = [z.astype('float32') for z in (x, y)]\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8a08392",
   "metadata": {},
   "source": [
    "Now we generate some data to train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1109b1d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "x, y = generate_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35069a6a",
   "metadata": {},
   "source": [
    "Here’s a plot of the training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1b1ffed",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.scatter(x, y)\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('y')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25da1785",
   "metadata": {},
   "source": [
    "We’ll also use data from the same process for cross-validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71b8bb98",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "x_validate, y_validate = generate_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cbfceea",
   "metadata": {},
   "source": [
    "## Models\n",
    "\n",
    "We supply functions to build two types of models.\n",
    "\n",
    "The first implements linear regression.\n",
    "\n",
    "This is achieved by constructing a neural network with just one layer, that maps\n",
    "to a single dimension (since the prediction is real-valued).\n",
    "\n",
    "The input `model` will be an instance of `keras.Sequential`, which is used to\n",
    "group a stack of layers into a single prediction model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3ea78da",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def build_regression_model(model):\n",
    "    model.add(Dense(units=1))\n",
    "    model.compile(optimizer=keras.optimizers.SGD(), \n",
    "                  loss='mean_squared_error')\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "529e2c37",
   "metadata": {},
   "source": [
    "In the function above you can see that we use stochastic gradient descent to\n",
    "train the model, and that the loss is mean squared error (MSE).\n",
    "\n",
    "MSE is the standard loss function for ordinary least squares regression.\n",
    "\n",
    "The second function creates a dense (i.e., fully connected) neural network with\n",
    "3 hidden layers, where each hidden layer maps to a k-dimensional output space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "826785da",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def build_nn_model(model, k=10, activation_function='tanh'):\n",
    "    # Construct network\n",
    "    model.add(Dense(units=k, activation=activation_function))\n",
    "    model.add(Dense(units=k, activation=activation_function))\n",
    "    model.add(Dense(units=k, activation=activation_function))\n",
    "    model.add(Dense(1))\n",
    "    # Embed training configurations\n",
    "    model.compile(optimizer=keras.optimizers.SGD(), \n",
    "                  loss='mean_squared_error')\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "185a39af",
   "metadata": {},
   "source": [
    "The following function will be used to plot the MSE of the model during the\n",
    "training process.\n",
    "\n",
    "Initially the MSE will be relatively high, but it should fall at each iteration,\n",
    "as the parameters are adjusted to better fit the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f942ab4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_loss_history(training_history, ax):\n",
    "    ax.plot(training_history.epoch, \n",
    "            np.array(training_history.history['loss']), \n",
    "            label='training loss')\n",
    "    ax.plot(training_history.epoch, \n",
    "            np.array(training_history.history['val_loss']),\n",
    "            label='validation loss')\n",
    "    ax.set_xlabel('Epoch')\n",
    "    ax.set_ylabel('Loss (Mean squared error)')\n",
    "    ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7023508e",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "Now let’s go ahead and train our  models."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2fb9560",
   "metadata": {},
   "source": [
    "### Linear regression\n",
    "\n",
    "We’ll start with linear regression.\n",
    "\n",
    "First we create a `Model` instance using `Sequential()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e64ca7bf",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "model = Sequential()\n",
    "regression_model = build_regression_model(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7d05198",
   "metadata": {},
   "source": [
    "Now we train the model using the training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59d35e28",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "training_history = regression_model.fit(\n",
    "    x, y, batch_size=x.shape[0], verbose=0,\n",
    "    epochs=4000, validation_data=(x_validate, y_validate))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "911eefa5",
   "metadata": {},
   "source": [
    "Let’s have a look at the evolution of MSE as the model is trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fdfb477",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "plot_loss_history(training_history, ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb017aff",
   "metadata": {},
   "source": [
    "Let’s print the final MSE on the cross-validation data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc3e8b0b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"Testing loss on the validation set.\")\n",
    "regression_model.evaluate(x_validate, y_validate, verbose=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b888eba3",
   "metadata": {},
   "source": [
    "Here’s our output predictions on the cross-validation data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3df3592b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "y_predict = regression_model.predict(x_validate, verbose=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08d7472c",
   "metadata": {},
   "source": [
    "We use the following function to plot our predictions along with the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc17652c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_results(x, y, y_predict, ax):\n",
    "    ax.scatter(x, y)\n",
    "    ax.plot(x, y_predict, label=\"fitted model\", color='black')\n",
    "    ax.set_xlabel('x')\n",
    "    ax.set_ylabel('y')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d5930f8",
   "metadata": {},
   "source": [
    "Let’s now call the function on the cross-validation data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dade5f4c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "plot_results(x_validate, y_validate, y_predict, ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91e9363a",
   "metadata": {},
   "source": [
    "### Deep learning\n",
    "\n",
    "Now let’s switch to a neural network with multiple layers.\n",
    "\n",
    "We implement the same steps as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5f2cf2c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "model = Sequential()\n",
    "nn_model = build_nn_model(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "680e066a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "training_history = nn_model.fit(\n",
    "    x, y, batch_size=x.shape[0], verbose=0,\n",
    "    epochs=4000, validation_data=(x_validate, y_validate))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaf211cf",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "plot_loss_history(training_history, ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "beb06567",
   "metadata": {},
   "source": [
    "Here’s the final MSE for the deep learning model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7841e7ab",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"Testing loss on the validation set.\")\n",
    "nn_model.evaluate(x_validate, y_validate, verbose=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3d4b958",
   "metadata": {},
   "source": [
    "You will notice that this loss is much lower than the one we achieved with\n",
    "linear regression, suggesting a better fit.\n",
    "\n",
    "To confirm this, let’s look at the fitted function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75311a01",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "y_predict = nn_model.predict(x_validate, verbose=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26ec614b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_results(x, y, y_predict, ax):\n",
    "    ax.scatter(x, y)\n",
    "    ax.plot(x, y_predict, label=\"fitted model\", color='black')\n",
    "    ax.set_xlabel('x')\n",
    "    ax.set_ylabel('y')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38fb81ae",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "plot_results(x_validate, y_validate, y_predict, ax)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "date": 1732064317.8912125,
  "filename": "keras.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Simple Neural Network Regression with Keras and JAX"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}