TensorFlowIntro.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TensorFlow\n",
    "\n",
    "We don't have to code up back propagation for every possible function or neural network architecture that we want to fit. There are lots of libraries targeted towards machine learning that make this task easy and computationally efficient. One of the most popular libraries is [TensorFlow](https://www.tensorflow.org/). It was developed by Google Brain and is now open source under the Apache License 2.0.\n",
    "\n",
    "(Other popular choices in 2023 are [PyTorch](https://pytorch.org/) and [JAX](https://jax.readthedocs.io/))\n",
    "\n",
    "The workflow consists of building a computational graph where \"operations\" act on \"tensors\" that can be automatically differentiated. Starting from tensorflow version 2 the operations are by default executed \"eagerly\" such that one can work with tensors in a similar way as with numpy arrays and typically does not have to worry about building the graph.\n",
    "\n",
    "The TensorFlow website contains a much more [detailed introduction](https://www.tensorflow.org/guide/low_level_intro) if you want to learn more."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numpy-like syntax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tensors can be created via `tf.constant` from python lists or numpy arrays. Similar to numpy arrays, they have a `shape` and a `dtype`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.constant([1, 2, 3], dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.constant(np.array([1, 2, 3]), dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.constant([[1, 2], [3, 4], [5, 6]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.constant([[[1, 2, 3, 4],\n",
    "              [5, 6, 7, 8],\n",
    "              [9, 10, 11, 12]],\n",
    "             [[13, 14, 15, 16],\n",
    "              [17, 18, 19, 20],\n",
    "              [21, 22, 23, 24]]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are also convenience functions, e.g. to create equidistant or random values and all sorts of mathematical functions that represent operations on tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.random.uniform((10, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = tf.linspace(0., 2.*np.pi, 10)\n",
    "t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "2 * t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.sin(t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tensors can be plotted like numpy arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(t, tf.sin(t))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or explicitely converted via `.numpy()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.sin(t).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Auto differentiation\n",
    "The real power comes from tracing operations that allows automatic backpropagation to calculate gradients. This can be done using `tf.GradientTape`. By default the gradients w.r.t. tensors (constants) are not recorded, but only for `tf.Variable`. A `tf.Variable` represents a mutable state - this makes sense, since in many cases we want to modify the values on which we calculate gradients (e.g. training a neural network)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = tf.Variable(tf.linspace(0., 2.*np.pi, 100))\n",
    "t"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now calculate the derivative of the `sin` function w.r.t. `t` using `tf.GradientTape` in a context manager"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.GradientTape() as tape:\n",
    "    f = tf.sin(t)\n",
    "df = tape.gradient(f, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: for plotting tf.Variable one always has to explicitely convert via .numpy()\n",
    "# (not nescessary for Tensors/tf.constant)\n",
    "plt.plot(t.numpy(), f, label=\"sin(t)\")\n",
    "plt.plot(t.numpy(), df, label=\"sin'(t)\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To calculate gradients w.r.t. Tensors (`tf.constant`) instead of `tf.Variable`, use `tape.watch`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_const = tf.linspace(0., 2.*np.pi, 100)\n",
    "with tf.GradientTape() as tape:\n",
    "    tape.watch(t_const)\n",
    "    f = tf.sin(t_const)\n",
    "plt.plot(t_const, f, label=\"sin(t)\")\n",
    "plt.plot(t_const, tape.gradient(f, t_const), label=\"sin'(t)\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The computation of the gradient can also be recorded and we can calculate the gradient of the gradient to get the second derivative."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.GradientTape() as tape0:\n",
    "    with tf.GradientTape() as tape1:\n",
    "        f = tf.sin(t)\n",
    "    df = tape1.gradient(f, t)\n",
    "ddf = tape0.gradient(df, t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The two gradient tapes are nescessary since tensorflow by default only allows one gradient to be calculated from a tape. If recording gradients themselves to the tape is intended one has to pass `persistent=True` - so the following works as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.GradientTape(persistent=True) as tape:\n",
    "    f = tf.sin(t)\n",
    "    # this is inside the with block, so the gradient itself will also be recorded to the gradient tape\n",
    "    df = tape.gradient(f, t)\n",
    "# now we can calculate the gradient of the gradient\n",
    "ddf_alternative = tape.gradient(df, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(t.numpy(), f.numpy(), label=\"sin(t)\")\n",
    "plt.plot(t.numpy(), df.numpy(), label=\"sin'(t)\")\n",
    "plt.plot(t.numpy(), ddf.numpy(), label=\"sin''(t)\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Manually build a NN in TensorFlow\n",
    "\n",
    "Let's build a 1-hidden-layer NN, similar to what we did in [NNFromScratchNumpy.ipynb](extra_material/NNFromScratchNumpy.ipynb) now with TensorFlow."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we define the weights and biases for the hidden layer and the output layer via `tf.Variable`. We initialize the weights randomly (normal distribution) and the biases to 0. We will again use the convention with column vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hidden layer parameters\n",
    "W = tf.Variable(tf.random.normal((16, 2)), name=\"W\")\n",
    "b = tf.Variable(tf.zeros((16, 1)))\n",
    "print(W)\n",
    "print(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# output layer parameters\n",
    "W2 = tf.Variable(tf.random.normal((1, 16)))\n",
    "b2 = tf.Variable(tf.zeros((1, 1)))\n",
    "print(W2)\n",
    "print(b2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's propagate some inputs through the neural network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp = tf.random.normal((10, 2, 1))\n",
    "inp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output of the first hidden layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = tf.add(tf.matmul(W, inp), b)\n",
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = tf.nn.relu(z)\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z2 = tf.add(tf.matmul(W2, a), b2)\n",
    "z2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will skip the activation function, since we will use a loss function that already applies the sigmoid transformation. This is numerically more stable.\n",
    "\n",
    "But first, we need to define the labels. For this experiment, let's choose them randomly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = tf.Variable(np.random.randint(0, 2, size=inp.shape[0]).reshape(-1, 1, 1), dtype=tf.float32)\n",
    "y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the binary cross entropy with a sigmoid transformation of the input values that don't have the sigmoid applied already."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = tf.nn.sigmoid_cross_entropy_with_logits(labels=y, logits=z2)\n",
    "L"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now - **and this is the whole point of this tutorial** - to get the gradients w.r.t. all parameters, we can record the parameters to a `tf.GradientTape` and get the gradients w.r.t. all parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward_NN(inp):\n",
    "    z = tf.add(tf.matmul(W, inp), b)\n",
    "    a = tf.nn.relu(z)\n",
    "    z2 = tf.add(tf.matmul(W2, a), b2)\n",
    "    return tf.nn.sigmoid_cross_entropy_with_logits(labels=y, logits=z2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.GradientTape() as tape:\n",
    "    L = forward_NN(inp)\n",
    "parameters = dict(W=W, b=b, W2=W2, b2=b2)\n",
    "# the gradient will have the same structure (dict, tuple) as the parameters\n",
    "grad_NN = tape.gradient(L, parameters)\n",
    "grad_NN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inspect the computation graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another way to trace computation graphs (other than `tf.GradientTape`) is by wrapping the operations into a `tf.function`. We won't go into detail here - have a look at the [`tf.function` tutorials](https://www.tensorflow.org/guide/intro_to_graphs) for further information.\n",
    "\n",
    "For illustration we use it here to display the computation graph with the `tensorboard` extension for Jupyter notebooks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from https://www.tensorflow.org/tensorboard/graphs\n",
    "\n",
    "from datetime import datetime\n",
    "logdir = f\"logs/{datetime.now().strftime('%Y%m%d-%H%M%S')}\"\n",
    "writer = tf.summary.create_file_writer(logdir)\n",
    "\n",
    "# Bracket the function call with\n",
    "# tf.summary.trace_on() and tf.summary.trace_export().\n",
    "tf.summary.trace_on(graph=True, profiler=True)\n",
    "# Call only one tf.function when tracing.\n",
    "L = tf.function(forward_NN)(inp)\n",
    "with writer.as_default():\n",
    "    tf.summary.trace_export(\n",
    "          name=\"my_func_trace\",\n",
    "          step=0,\n",
    "          profiler_outdir=logdir\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you run this notebook remotely at CIP you need to forward the port for the tensorboard process (6006), e.g. via (in a terminal on your local machine)\n",
    "```\n",
    "ssh -L 6006:localhost:6006 <your-username>@<your-cip-host>\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!hostname"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%tensorboard --logdir $logdir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# kill the tensorboard process\n",
    "!sleep 3\n",
    "!killall tensorboard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare to NNFromScratch\n",
    "\n",
    "Let's see if we can reproduce that with the formulas we used in [NNFromScratchNumpy.ipynb](extra_material/NNFromScratchNumpy.ipynb)\n",
    "\n",
    "Here a copy paste of the relevant functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sigmoid(Z):\n",
    "    return 1/(1+np.exp(-Z))\n",
    "\n",
    "def relu(Z):\n",
    "    return np.maximum(0,Z)\n",
    "\n",
    "def sigmoid_derivative(Z):\n",
    "    sig = sigmoid(Z)\n",
    "    return sig * (1 - sig)\n",
    "\n",
    "def relu_derivative(Z):\n",
    "    dZ = (Z >= 0)\n",
    "    return dZ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def single_layer_forward_propagation(A_prev, W_curr, b_curr, activation=\"relu\"):\n",
    "    Z_curr = np.matmul(W_curr, A_prev) + b_curr\n",
    "    \n",
    "    if activation == \"relu\":\n",
    "        activation_func = relu\n",
    "    elif activation == \"sigmoid\":\n",
    "        activation_func = sigmoid\n",
    "    else:\n",
    "        raise Exception('Non-supported activation function')\n",
    "        \n",
    "    return activation_func(Z_curr), Z_curr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def single_layer_backward_propagation(dA_curr, W_curr, b_curr, Z_curr, A_prev, activation=\"relu\"):\n",
    "    \n",
    "    if activation == \"relu\":\n",
    "        derivative_activation_func = relu_derivative\n",
    "    elif activation == \"sigmoid\":\n",
    "        derivative_activation_func = sigmoid_derivative\n",
    "    else:\n",
    "        raise Exception('Non-supported activation function')\n",
    "            \n",
    "    dZ_curr = dA_curr * derivative_activation_func(Z_curr)\n",
    "    dW_curr = np.matmul(\n",
    "        dZ_curr,\n",
    "        # need to transpose only the last 2 dimensions, \n",
    "        # since the first dimension is the training example index\n",
    "        np.transpose(A_prev, (0, 2, 1))\n",
    "    )\n",
    "    db_curr = dZ_curr\n",
    "    dA_prev = np.matmul(W_curr.T, dZ_curr)\n",
    "\n",
    "    return dA_prev, dW_curr, db_curr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_loss_value(Y_hat, Y):\n",
    "    return - np.mean(Y * np.log(Y_hat) + (1 - Y) * np.log(1 - Y_hat))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_loss_derivative(Y_hat, Y):\n",
    "    return - (np.divide(Y, Y_hat) - np.divide(1 - Y, 1 - Y_hat));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First let's calculate the forward pass.\n",
    "\n",
    "First, lets store the initialized values of the NN parameters, inputs and labels in python variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inp_val, W_val, b_val, W2_val, b2_val, y_val = inp.numpy(), W.numpy(), b.numpy(), W2.numpy(), b2.numpy(), y.numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then run our manual `numpy` forward propagation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_val, z_val = single_layer_forward_propagation(inp_val, W_val, b_val)\n",
    "print(a_val[0].ravel())\n",
    "print(z_val[0].ravel())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "compared to what `tf` gave us"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(a[0].numpy().ravel())\n",
    "print(z[0].numpy().ravel())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a2_val, z2_val = single_layer_forward_propagation(\n",
    "    a_val, W2_val, b2_val, activation=\"sigmoid\"\n",
    ")\n",
    "print(a2_val.ravel())\n",
    "print(z2_val.ravel())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For `tf`, we don't have `a2` because we used a definition of the loss function where the sigmoid activation is already included. But we have `z2`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z2.numpy().ravel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great, we implemented the forward pass correctly, so now lets do the backward pass and check if we get the same gradients like Tensorflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dL = get_loss_derivative(a2_val, y_val)\n",
    "dL.ravel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Propagate back into the output layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da, dW2, db2 = single_layer_backward_propagation(\n",
    "    dL, W2_val, b2_val, z2_val, a_val, activation=\"sigmoid\"\n",
    ")\n",
    "print(np.sum(dW2, axis=0).ravel())\n",
    "print(np.sum(db2, axis=0).ravel())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(grad_NN[\"W2\"].numpy().ravel())\n",
    "print(grad_NN[\"b2\"].numpy().ravel())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And from there into the hidden layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dinp, dW, db = single_layer_backward_propagation(\n",
    "    da, W_val, b_val, z_val, inp_val, activation=\"relu\"\n",
    ")\n",
    "print(np.sum(dW, axis=0))\n",
    "print(np.sum(db, axis=0).ravel())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(grad_NN[\"W\"].numpy())\n",
    "print(grad_NN[\"b\"].numpy().ravel())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! Tensorflow does the same thing we attempted to do before."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.0"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}