{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercise: Training a digit-classification neural network on the MNIST dataset using Keras\n",
    "\n",
    "This example is from [Stefan Wunsch (CERN IML TensoFlow and Keras workshop)](https://github.com/stwunsch/iml_tensorflow_keras_workshop). See also the example on the [Keras website](https://keras.io/examples/vision/mnist_convnet/).\n",
    "\n",
    "The MNIST dataset is one of the most popular benchmark-datasets in modern machine learning. The dataset consists of 70000 images of handwritten digits and associated labels which can be used to train neural network performing image classification.\n",
    "\n",
    "The following program presents the basic workflow of Keras showing the most import details of the API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "np.random.seed(1234)\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download the dataset\n",
    "\n",
    "The code below downloads the dataset and performs a scaling of the pixel-values of the images. Because the images are encoded with 8-bit unsigned int values, we scale these values to floating-point values in the range `[0, 1)` so that the inputs match the activation of the neurons better."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.datasets import mnist\n",
    "from tensorflow.keras.utils import to_categorical\n",
    "\n",
    "# Download dataset\n",
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "\n",
    "# The data is loaded as flat array with 784 entries (28x28),\n",
    "# we need to reshape it into an array with shape:\n",
    "# (num_images, pixels_row, pixels_column, color channels)\n",
    "x_train = x_train.reshape(x_train.shape[0], 28, 28, 1)\n",
    "x_test = x_test.reshape(x_test.shape[0], 28, 28, 1)\n",
    "\n",
    "# Convert digits to one-hot vectors, e.g.,\n",
    "# 2 -> [0 0 1 0 0 0 0 0 0 0]\n",
    "# 0 -> [1 0 0 0 0 0 0 0 0 0]\n",
    "# 9 -> [0 0 0 0 0 0 0 0 0 1]\n",
    "y_train = to_categorical(y_train, 10)\n",
    "y_test = to_categorical(y_test, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Addtionally, we store some example images to disk to show later on the inference part of the Keras API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 6344 characters]",
      "text/plain": [
       "<Figure size 1200x200 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import png\n",
    "\n",
    "num_examples = 6\n",
    "# offset = 100\n",
    "offset = 200\n",
    "\n",
    "plt.figure(figsize=(num_examples*2, 2))\n",
    "for i in range(num_examples):\n",
    "    plt.subplot(1, num_examples, i+1)\n",
    "    plt.axis('off')\n",
    "    # example = np.squeeze(np.array(x_test[offset+i]*255).astype(\"uint8\"))\n",
    "    example = np.squeeze(np.array(x_test[offset+i]).astype(\"uint8\"))\n",
    "    plt.imshow(example, cmap=\"gray\")\n",
    "    w = png.Writer(28, 28, greyscale=True)\n",
    "    w.write(open(\"mnist_example_{}.png\".format(i+1), 'wb'), example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model / data parameters\n",
    "num_classes = 10\n",
    "input_shape = (28, 28, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the model\n",
    "\n",
    "The model definition in Keras can be done using the `Sequential` or the functional API. Shown here is the `Sequential` API allowing to stack neural network layers on top of each other, which is feasible for most neural network models. In contrast, the functional API would allow to have multiple inputs and outputs for a maximum of flexibility to build your custom model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      " Layer (type)                Output Shape              Param #   \n",
      "=================================================================\n",
      " conv2d (Conv2D)             (None, 26, 26, 8)         80        \n",
      "                                                                 \n",
      " max_pooling2d (MaxPooling2D  (None, 13, 13, 8)        0         \n",
      " )                                                               \n",
      "                                                                 \n",
      " flatten (Flatten)           (None, 1352)              0         \n",
      "                                                                 \n",
      " dense (Dense)               (None, 16)                21648     \n",
      "                                                                 \n",
      " dense_1 (Dense)             (None, 10)                170       \n",
      "                                                                 \n",
      "=================================================================\n",
      "Total params: 21,898\n",
      "Trainable params: 21,898\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.layers import Dense, Flatten, MaxPooling2D, Conv2D, Input, Dropout\n",
    "\n",
    "# conv layer with 8 3x3 filters\n",
    "model = Sequential(\n",
    "    [\n",
    "        Input(shape=input_shape),\n",
    "        Conv2D(8, kernel_size=(3, 3), activation=\"relu\"),\n",
    "        MaxPooling2D(pool_size=(2, 2)),\n",
    "        Flatten(),\n",
    "        Dense(16, activation=\"relu\"),\n",
    "        Dense(num_classes, activation=\"softmax\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compile the model\n",
    "\n",
    "Using Keras, you have to `compile` a model, which means adding the loss function, the optimizer algorithm and validation metrics to your training setup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile(loss=\"categorical_crossentropy\",\n",
    "        optimizer=\"adam\",\n",
    "        metrics=[\"accuracy\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model\n",
    "\n",
    "The cell below shows the training procedure of Keras using the `model.fit(...)` method. Besides typical options such as `batch_size` and `epochs`, which control the number of gradient steps of your training, Keras allows to use callbacks during training.\n",
    "\n",
    "Callbacks are methods, which are called during training to perform tasks such as saving checkpoints of the model (`ModelCheckpoint`) or stop the training early if a convergence criteria is met (`EarlyStopping`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/50\n",
      " 12/150 [=>............................] - ETA: 0s - loss: 17.0395 - accuracy: 0.0942 "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-04-14 11:48:41.118317: W tensorflow/tsl/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "140/150 [===========================>..] - ETA: 0s - loss: 3.4982 - accuracy: 0.1353\n",
      "Epoch 1: val_loss improved from inf to 1.84449, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 8ms/step - loss: 3.3907 - accuracy: 0.1458 - val_loss: 1.8445 - val_accuracy: 0.3156\n",
      "Epoch 2/50\n",
      "148/150 [============================>.] - ETA: 0s - loss: 1.6724 - accuracy: 0.3665\n",
      "Epoch 2: val_loss improved from 1.84449 to 1.58089, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 1.6707 - accuracy: 0.3670 - val_loss: 1.5809 - val_accuracy: 0.3983\n",
      "Epoch 3/50\n",
      "148/150 [============================>.] - ETA: 0s - loss: 1.4790 - accuracy: 0.4209\n",
      "Epoch 3: val_loss improved from 1.58089 to 1.33352, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 1.4755 - accuracy: 0.4219 - val_loss: 1.3335 - val_accuracy: 0.5253\n",
      "Epoch 4/50\n",
      "146/150 [============================>.] - ETA: 0s - loss: 1.2029 - accuracy: 0.5586\n",
      "Epoch 4: val_loss improved from 1.33352 to 1.18135, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 1.2032 - accuracy: 0.5593 - val_loss: 1.1813 - val_accuracy: 0.5806\n",
      "Epoch 5/50\n",
      "140/150 [===========================>..] - ETA: 0s - loss: 1.0944 - accuracy: 0.5868\n",
      "Epoch 5: val_loss improved from 1.18135 to 1.09381, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 1.0921 - accuracy: 0.5873 - val_loss: 1.0938 - val_accuracy: 0.5917\n",
      "Epoch 6/50\n",
      "143/150 [===========================>..] - ETA: 0s - loss: 1.0173 - accuracy: 0.6051\n",
      "Epoch 6: val_loss improved from 1.09381 to 1.04734, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 1.0192 - accuracy: 0.6051 - val_loss: 1.0473 - val_accuracy: 0.6206\n",
      "Epoch 7/50\n",
      "143/150 [===========================>..] - ETA: 0s - loss: 0.9230 - accuracy: 0.6555\n",
      "Epoch 7: val_loss improved from 1.04734 to 0.89113, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.9215 - accuracy: 0.6564 - val_loss: 0.8911 - val_accuracy: 0.6717\n",
      "Epoch 8/50\n",
      "145/150 [============================>.] - ETA: 0s - loss: 0.7411 - accuracy: 0.7521\n",
      "Epoch 8: val_loss improved from 0.89113 to 0.70894, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.7397 - accuracy: 0.7534 - val_loss: 0.7089 - val_accuracy: 0.7917\n",
      "Epoch 9/50\n",
      "150/150 [==============================] - ETA: 0s - loss: 0.6071 - accuracy: 0.8192\n",
      "Epoch 9: val_loss improved from 0.70894 to 0.61332, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 8ms/step - loss: 0.6071 - accuracy: 0.8192 - val_loss: 0.6133 - val_accuracy: 0.8285\n",
      "Epoch 10/50\n",
      "147/150 [============================>.] - ETA: 0s - loss: 0.5089 - accuracy: 0.8581\n",
      "Epoch 10: val_loss improved from 0.61332 to 0.49433, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.5086 - accuracy: 0.8583 - val_loss: 0.4943 - val_accuracy: 0.8716\n",
      "Epoch 11/50\n",
      "147/150 [============================>.] - ETA: 0s - loss: 0.3716 - accuracy: 0.9059\n",
      "Epoch 11: val_loss improved from 0.49433 to 0.36685, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.3701 - accuracy: 0.9060 - val_loss: 0.3669 - val_accuracy: 0.9154\n",
      "Epoch 12/50\n",
      "149/150 [============================>.] - ETA: 0s - loss: 0.2866 - accuracy: 0.9238\n",
      "Epoch 12: val_loss improved from 0.36685 to 0.32505, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.2869 - accuracy: 0.9238 - val_loss: 0.3250 - val_accuracy: 0.9172\n",
      "Epoch 13/50\n",
      "146/150 [============================>.] - ETA: 0s - loss: 0.2295 - accuracy: 0.9359\n",
      "Epoch 13: val_loss improved from 0.32505 to 0.27026, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.2292 - accuracy: 0.9359 - val_loss: 0.2703 - val_accuracy: 0.9292\n",
      "Epoch 14/50\n",
      "146/150 [============================>.] - ETA: 0s - loss: 0.1937 - accuracy: 0.9450\n",
      "Epoch 14: val_loss improved from 0.27026 to 0.24167, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.1949 - accuracy: 0.9447 - val_loss: 0.2417 - val_accuracy: 0.9374\n",
      "Epoch 15/50\n",
      "142/150 [===========================>..] - ETA: 0s - loss: 0.1748 - accuracy: 0.9496\n",
      "Epoch 15: val_loss did not improve from 0.24167\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.1737 - accuracy: 0.9501 - val_loss: 0.2418 - val_accuracy: 0.9384\n",
      "Epoch 16/50\n",
      "143/150 [===========================>..] - ETA: 0s - loss: 0.1605 - accuracy: 0.9535\n",
      "Epoch 16: val_loss improved from 0.24167 to 0.23019, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.1607 - accuracy: 0.9536 - val_loss: 0.2302 - val_accuracy: 0.9406\n",
      "Epoch 17/50\n",
      "142/150 [===========================>..] - ETA: 0s - loss: 0.1472 - accuracy: 0.9569\n",
      "Epoch 17: val_loss improved from 0.23019 to 0.22133, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.1466 - accuracy: 0.9571 - val_loss: 0.2213 - val_accuracy: 0.9423\n",
      "Epoch 18/50\n",
      "149/150 [============================>.] - ETA: 0s - loss: 0.1381 - accuracy: 0.9596\n",
      "Epoch 18: val_loss improved from 0.22133 to 0.20980, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.1378 - accuracy: 0.9596 - val_loss: 0.2098 - val_accuracy: 0.9475\n",
      "Epoch 19/50\n",
      "145/150 [============================>.] - ETA: 0s - loss: 0.1266 - accuracy: 0.9632\n",
      "Epoch 19: val_loss improved from 0.20980 to 0.20645, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.1264 - accuracy: 0.9631 - val_loss: 0.2065 - val_accuracy: 0.9483\n",
      "Epoch 20/50\n",
      "141/150 [===========================>..] - ETA: 0s - loss: 0.1222 - accuracy: 0.9652\n",
      "Epoch 20: val_loss did not improve from 0.20645\n",
      "150/150 [==============================] - 1s 8ms/step - loss: 0.1222 - accuracy: 0.9652 - val_loss: 0.2111 - val_accuracy: 0.9489\n",
      "Epoch 21/50\n",
      "146/150 [============================>.] - ETA: 0s - loss: 0.1142 - accuracy: 0.9661\n",
      "Epoch 21: val_loss improved from 0.20645 to 0.19401, saving model to mnist_keras_model.h5\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.1136 - accuracy: 0.9664 - val_loss: 0.1940 - val_accuracy: 0.9506\n",
      "Epoch 22/50\n",
      "142/150 [===========================>..] - ETA: 0s - loss: 0.1114 - accuracy: 0.9683\n",
      "Epoch 22: val_loss did not improve from 0.19401\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.1105 - accuracy: 0.9684 - val_loss: 0.1999 - val_accuracy: 0.9514\n",
      "Epoch 23/50\n",
      "142/150 [===========================>..] - ETA: 0s - loss: 0.1030 - accuracy: 0.9696\n",
      "Epoch 23: val_loss did not improve from 0.19401\n",
      "150/150 [==============================] - 1s 7ms/step - loss: 0.1037 - accuracy: 0.9692 - val_loss: 0.1991 - val_accuracy: 0.9509\n"
     ]
    }
   ],
   "source": [
    "from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping\n",
    "\n",
    "checkpoint = ModelCheckpoint(\n",
    "            filepath=\"mnist_keras_model.h5\",\n",
    "            save_best_only=True,\n",
    "            verbose=1)\n",
    "early_stopping = EarlyStopping(patience=2)\n",
    "\n",
    "history = model.fit(x_train, y_train, # Training data\n",
    "            batch_size=200, # Batch size\n",
    "            epochs=50, # Maximum number of training epochs\n",
    "            validation_split=0.5, # Use 50% of the train dataset for validation\n",
    "            callbacks=[checkpoint, early_stopping]) # Register callbacks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### a) Plot training and validation loss as well as training and validation accurace as a function of the number of epochs "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 68624 characters]",
      "text/plain": [
       "<Figure size 1200x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "epochs = range(1, len(history.history[\"loss\"])+1)\n",
    "plt.figure(figsize=(12,5))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(epochs, history.history[\"loss\"], label=\"Training loss\")\n",
    "plt.plot(epochs, history.history[\"val_loss\"], label=\"Validation loss\")\n",
    "plt.legend(fontsize=15), plt.xlabel(\"Epochs\", fontsize=15), plt.ylabel(\"Loss\", fontsize=15)\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(epochs, history.history[\"accuracy\"], label=\"Training accuracy\")\n",
    "plt.plot(epochs, history.history[\"val_accuracy\"], label=\"Validation accuracy\")\n",
    "plt.legend(fontsize=15), plt.xlabel(\"Epochs\", fontsize=15), plt.ylabel(\"Accuracy\", fontsize=15);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### b) Determine the accuracy of the fully trained model\n",
    "\n",
    "The prediction of unseen data is performed using the `model.predict(inputs)` call. Below, a basic test of the model is done by calculating the accuracy on the test dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "313/313 [==============================] - 0s 621us/step\n",
      "Test accuracy: 0.9522\n"
     ]
    }
   ],
   "source": [
    "### Your code here ###\n",
    "\n",
    "# Get predictions on test dataset\n",
    "y_pred = model.predict(x_test)\n",
    "\n",
    "# Compare predictions with ground truth\n",
    "test_accuracy = np.sum(\n",
    "        np.argmax(y_test, axis=1)==np.argmax(y_pred, axis=1))/float(x_test.shape[0])\n",
    "\n",
    "print(\"Test accuracy: {}\".format(test_accuracy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "313/313 [==============================] - 0s 788us/step - loss: 0.1853 - accuracy: 0.9522\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[0.1852734535932541, 0.9521999955177307]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.evaluate(x_test, y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
