tensorflow2_metrics_summaries.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "tensorflow2_metrics_summaries.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/adventuresinML/adventures-in-ml-code/blob/master/tensorflow2_metrics_summaries.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v63fE79TzddR",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 88
        },
        "outputId": "e93de207-3c62-48c7-e530-87080158ce6f"
      },
      "source": [
        "!pip install -q tensorflow==2.0.0-alpha0"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\u001b[K     |████████████████████████████████| 79.9MB 72.2MB/s \n",
            "\u001b[K     |████████████████████████████████| 61kB 24.7MB/s \n",
            "\u001b[K     |████████████████████████████████| 3.0MB 44.4MB/s \n",
            "\u001b[K     |████████████████████████████████| 419kB 47.2MB/s \n",
            "\u001b[?25h"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sT-z6b6mwg_s",
        "colab_type": "code",
        "outputId": "908ba78c-4188-4508-fa49-097cfc24b879",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 248
        }
      },
      "source": [
        "!wget https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip\n",
        "!unzip ngrok-stable-linux-amd64.zip\n",
        "LOG_DIR = './log'\n",
        "get_ipython().system_raw(\n",
        "    'tensorboard --logdir {} --host 0.0.0.0 --port 6006 &'\n",
        "    .format(LOG_DIR)\n",
        ")"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2019-05-05 22:12:13--  https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip\n",
            "Resolving bin.equinox.io (bin.equinox.io)... 52.202.60.111, 34.196.237.103, 52.203.102.189, ...\n",
            "Connecting to bin.equinox.io (bin.equinox.io)|52.202.60.111|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 14991793 (14M) [application/octet-stream]\n",
            "Saving to: ‘ngrok-stable-linux-amd64.zip’\n",
            "\n",
            "ngrok-stable-linux- 100%[===================>]  14.30M  5.98MB/s    in 2.4s    \n",
            "\n",
            "2019-05-05 22:12:16 (5.98 MB/s) - ‘ngrok-stable-linux-amd64.zip’ saved [14991793/14991793]\n",
            "\n",
            "Archive:  ngrok-stable-linux-amd64.zip\n",
            "  inflating: ngrok                   \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R9-0-P_L236s",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "import datetime as dt\n",
        "import pdb"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DvqbVJSbIfLT",
        "colab_type": "text"
      },
      "source": [
        "#### Simple metric example"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ts07TPPGIbcS",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "5b48b455-a681-4488-ad50-eb6546bb46a1"
      },
      "source": [
        "mean_metric = tf.keras.metrics.Mean()\n",
        "mean_metric.update_state(2.0)\n",
        "mean_metric.update_state(3.0)\n",
        "mean_metric.update_state(4.0)\n",
        "print(mean_metric.result().numpy())"
      ],
      "execution_count": 60,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "3.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jucs14c7I9r_",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "ae6b294c-eb02-441c-b95c-9c23a35dba97"
      },
      "source": [
        "mean_metric.reset_states()\n",
        "print(mean_metric.result().numpy())"
      ],
      "execution_count": 61,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "0.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jKLdSrhgJIFl",
        "colab_type": "text"
      },
      "source": [
        "#### MNIST CNN example"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PQE9WMrStSm5",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 52
        },
        "outputId": "d8256868-8d75-4361-a73c-8698bc721e74"
      },
      "source": [
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz\n",
            "11493376/11490434 [==============================] - 0s 0us/step\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kMSCPBIGy2T4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "BATCH_SIZE=64\n",
        "train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(BATCH_SIZE).shuffle(10000)\n",
        "train_dataset = train_dataset.map(lambda x, y: (tf.cast(x, tf.float32) / 255.0, y))\n",
        "train_dataset = train_dataset.map(lambda x, y: (tf.expand_dims(x, -1) / 255.0, y))\n",
        "train_dataset = train_dataset.repeat()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HpxL1shN1qsp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "valid_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(5000).shuffle(10000)\n",
        "valid_dataset = valid_dataset.map(lambda x, y: (tf.cast(x, tf.float32) / 255.0, y))\n",
        "valid_dataset = valid_dataset.map(lambda x, y: (tf.expand_dims(x, -1) / 255.0, y))\n",
        "valid_dataset = valid_dataset.repeat()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kf7KYLcatwuB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = tf.keras.Sequential()\n",
        "model.add(tf.keras.layers.Conv2D(32, 2, 1, activation='relu', input_shape=(28, 28, 1)))\n",
        "model.add(tf.keras.layers.MaxPool2D(2))\n",
        "model.add(tf.keras.layers.BatchNormalization())\n",
        "model.add(tf.keras.layers.Conv2D(32, 2, 1, activation='relu'))\n",
        "model.add(tf.keras.layers.MaxPool2D(2))\n",
        "model.add(tf.keras.layers.BatchNormalization())\n",
        "model.add(tf.keras.layers.Flatten())\n",
        "model.add(tf.keras.layers.Dense(10))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aEPUyxxpzJyk",
        "colab_type": "text"
      },
      "source": [
        "## Raw TensorFlow implementation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "70aYSk3_5jUk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "optimizer = tf.keras.optimizers.Adam()\n",
        "loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rGMFUsz44X5d",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(ds_train, optimizer, loss_fn, model, num_batches, log_freq=10):\n",
        "  avg_loss = tf.keras.metrics.Mean()\n",
        "  avg_acc = tf.keras.metrics.SparseCategoricalAccuracy()\n",
        "  batch_idx = 0\n",
        "  for batch_idx, (images, labels) in enumerate(ds_train):\n",
        "    images = tf.expand_dims(images, -1)\n",
        "    with tf.GradientTape() as tape:\n",
        "      logits = model(images)\n",
        "      loss_value = loss_fn(labels, logits)\n",
        "    grads = tape.gradient(loss_value, model.trainable_variables)\n",
        "    optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
        "    avg_loss.update_state(loss_value)\n",
        "    avg_acc.update_state(labels, logits)\n",
        "    if batch_idx % log_freq == 0:\n",
        "      print(f\"Batch {batch_idx}, average loss is {avg_loss.result().numpy()}, average accuracy is {avg_acc.result().numpy()}\")\n",
        "      tf.summary.scalar('loss', avg_loss.result(), step=optimizer.iterations)\n",
        "      tf.summary.scalar('acc', avg_acc.result(), step=optimizer.iterations)\n",
        "      avg_loss.reset_states()\n",
        "      avg_acc.reset_states()\n",
        "    if batch_idx > num_batches:\n",
        "      break"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qzmioFmDgpSF",
        "colab_type": "code",
        "outputId": "2feb93ef-1e55-4fa3-ac55-986986a19c62",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "get_ipython().system_raw('./ngrok http 6006 &')\n",
        "! curl -s http://localhost:4040/api/tunnels | python3 -c \\\n",
        "    \"import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])\""
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "https://e5478e54.ngrok.io\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SPdxcVVW1ITD",
        "colab_type": "code",
        "outputId": "49c03059-91bf-4679-aea1-5f8e149fe661",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1272
        }
      },
      "source": [
        "num_epochs = 10\n",
        "summary_writer = tf.summary.create_file_writer('./log/{}'.format(dt.datetime.now().strftime(\"%Y-%m-%d-%H-%M-%S\")))\n",
        "for i in range(num_epochs):\n",
        "  print(f\"Epoch {i + 1} of {num_epochs}\")\n",
        "  with summary_writer.as_default():\n",
        "    train(train_dataset, optimizer, loss_fn, model, 10000//BATCH_SIZE)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch 1 of 10\n",
            "Batch 0, average loss is 2.321280002593994, average accuracy is 0.078125\n",
            "Batch 10, average loss is 2.127411365509033, average accuracy is 0.30937498807907104\n",
            "Batch 20, average loss is 1.5143563747406006, average accuracy is 0.609375\n",
            "Batch 30, average loss is 0.9550114870071411, average accuracy is 0.6796875\n",
            "Batch 40, average loss is 0.6596224904060364, average accuracy is 0.7593749761581421\n",
            "Batch 50, average loss is 0.5405800342559814, average accuracy is 0.859375\n",
            "Batch 60, average loss is 0.4321937561035156, average accuracy is 0.8765624761581421\n",
            "Batch 70, average loss is 0.43382543325424194, average accuracy is 0.856249988079071\n",
            "Batch 80, average loss is 0.29824918508529663, average accuracy is 0.909375011920929\n",
            "Batch 90, average loss is 0.33607107400894165, average accuracy is 0.8890625238418579\n",
            "Batch 100, average loss is 0.34584954380989075, average accuracy is 0.890625\n",
            "Batch 110, average loss is 0.27703607082366943, average accuracy is 0.9203125238418579\n",
            "Batch 120, average loss is 0.2640536427497864, average accuracy is 0.9203125238418579\n",
            "Batch 130, average loss is 0.28403061628341675, average accuracy is 0.921875\n",
            "Batch 140, average loss is 0.27499693632125854, average accuracy is 0.9309210777282715\n",
            "Batch 150, average loss is 0.2746924161911011, average accuracy is 0.9312499761581421\n",
            "Epoch 2 of 10\n",
            "Batch 0, average loss is 0.2426912486553192, average accuracy is 0.9375\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-86-b68785b15f0d>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      4\u001b[0m   \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"Epoch {i + 1} of {num_epochs}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m   \u001b[0;32mwith\u001b[0m \u001b[0msummary_writer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mas_default\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m     \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_dataset\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mloss_fn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m10000\u001b[0m\u001b[0;34m//\u001b[0m\u001b[0mBATCH_SIZE\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m<ipython-input-79-dd6246ef270a>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(ds_train, optimizer, loss_fn, model, num_batches, log_freq)\u001b[0m\n\u001b[1;32m      9\u001b[0m       \u001b[0mlogits\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimages\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m       \u001b[0mloss_value\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mloss_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlabels\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlogits\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 11\u001b[0;31m     \u001b[0mgrads\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtape\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgradient\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss_value\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainable_variables\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     12\u001b[0m     \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mapply_gradients\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgrads\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainable_variables\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m     \u001b[0mavg_loss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mupdate_state\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss_value\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/backprop.py\u001b[0m in \u001b[0;36mgradient\u001b[0;34m(self, target, sources, output_gradients, unconnected_gradients)\u001b[0m\n\u001b[1;32m    954\u001b[0m         \u001b[0mflat_sources\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    955\u001b[0m         \u001b[0moutput_gradients\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0moutput_gradients\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 956\u001b[0;31m         unconnected_gradients=unconnected_gradients)\n\u001b[0m\u001b[1;32m    957\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    958\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_persistent\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/imperative_grad.py\u001b[0m in \u001b[0;36mimperative_grad\u001b[0;34m(tape, target, sources, output_gradients, unconnected_gradients)\u001b[0m\n\u001b[1;32m     70\u001b[0m       \u001b[0msources\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     71\u001b[0m       \u001b[0moutput_gradients\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 72\u001b[0;31m       compat.as_str(unconnected_gradients.value))\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/backprop.py\u001b[0m in \u001b[0;36m_gradient_function\u001b[0;34m(op_name, attr_tuple, num_inputs, inputs, outputs, out_grads)\u001b[0m\n\u001b[1;32m    131\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mnum_inputs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    132\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 133\u001b[0;31m   \u001b[0;32mreturn\u001b[0m \u001b[0mgrad_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmock_op\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mout_grads\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    134\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    135\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/nn_grad.py\u001b[0m in \u001b[0;36m_Conv2DGrad\u001b[0;34m(op, grad)\u001b[0m\n\u001b[1;32m    599\u001b[0m           \u001b[0mexplicit_paddings\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mexplicit_paddings\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    600\u001b[0m           \u001b[0muse_cudnn_on_gpu\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0muse_cudnn_on_gpu\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 601\u001b[0;31m           data_format=data_format)\n\u001b[0m\u001b[1;32m    602\u001b[0m   ]\n\u001b[1;32m    603\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/gen_nn_ops.py\u001b[0m in \u001b[0;36mconv2d_backprop_filter\u001b[0;34m(input, filter_sizes, out_backprop, strides, padding, use_cudnn_on_gpu, explicit_paddings, data_format, dilations, name)\u001b[0m\n\u001b[1;32m   1185\u001b[0m         \u001b[0mfilter_sizes\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_backprop\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"strides\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstrides\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"use_cudnn_on_gpu\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1186\u001b[0m         \u001b[0muse_cudnn_on_gpu\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"padding\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpadding\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"explicit_paddings\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1187\u001b[0;31m         explicit_paddings, \"data_format\", data_format, \"dilations\", dilations)\n\u001b[0m\u001b[1;32m   1188\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0m_result\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1189\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0m_core\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_FallbackException\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tKxruKgyzXD7",
        "colab_type": "text"
      },
      "source": [
        "## Keras training loop implementation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3qqvWya14BoV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = tf.keras.Sequential()\n",
        "model.add(tf.keras.layers.Conv2D(32, 2, 1, activation='relu', input_shape=(28, 28, 1)))\n",
        "model.add(tf.keras.layers.MaxPool2D(2))\n",
        "model.add(tf.keras.layers.BatchNormalization())\n",
        "model.add(tf.keras.layers.Conv2D(32, 2, 1, activation='relu'))\n",
        "model.add(tf.keras.layers.MaxPool2D(2))\n",
        "model.add(tf.keras.layers.BatchNormalization())\n",
        "model.add(tf.keras.layers.Flatten())\n",
        "model.add(tf.keras.layers.Dense(10))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_8uF7Dsx0vQw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class MetricLayer(tf.keras.layers.Layer):\n",
        "  def __init__(self, layer_to_log):\n",
        "    super(MetricLayer, self).__init__()\n",
        "    self.layer_to_log = layer_to_log\n",
        "    \n",
        "  def call(self, input):\n",
        "    self.add_metric(tf.keras.backend.std(self.layer_to_log.variables[0]),\n",
        "                    name=f'std_of_{self.layer_to_log.name}_kernel',\n",
        "                    aggregation='mean')\n",
        "    return input"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qB9MWMLD0ytT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "metric_model = tf.keras.Sequential()\n",
        "metric_model.add(model)\n",
        "metric_model.add(MetricLayer(model.layers[0]))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9oiLY4iy1FCY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "callbacks = [\n",
        "  # Write TensorBoard logs to `./logs` directory\n",
        "  tf.keras.callbacks.TensorBoard(log_dir='./log/{}'.format(dt.datetime.now().strftime(\"%Y-%m-%d-%H-%M-%S\")), update_freq='batch')\n",
        "]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3liSDPG40SRk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "metric_model.compile(optimizer=tf.optimizers.Adam(),\n",
        "                     loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "                     metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_zRexjVVeN5N",
        "colab_type": "code",
        "outputId": "40c97fc3-875c-4afa-e1b1-a44985ef8a04",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 232
        }
      },
      "source": [
        "metric_model.fit(train_dataset, steps_per_epoch=10000//BATCH_SIZE, epochs=5,\n",
        "                 validation_data=valid_dataset, validation_steps=5,\n",
        "                 callbacks=callbacks)"
      ],
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/5\n",
            "156/156 [==============================] - 12s 79ms/step - loss: 0.6518 - sparse_categorical_accuracy: 0.8202 - std_of_conv2d_2_kernel: 0.1359 - val_loss: 2.2652 - val_sparse_categorical_accuracy: 0.0888 - val_std_of_conv2d_2_kernel: 0.1454\n",
            "Epoch 2/5\n",
            "156/156 [==============================] - 16s 101ms/step - loss: 0.2318 - sparse_categorical_accuracy: 0.9311 - std_of_conv2d_2_kernel: 0.1523 - val_loss: 2.1332 - val_sparse_categorical_accuracy: 0.1647 - val_std_of_conv2d_2_kernel: 0.1585\n",
            "Epoch 3/5\n",
            "156/156 [==============================] - 16s 103ms/step - loss: 0.1493 - sparse_categorical_accuracy: 0.9578 - std_of_conv2d_2_kernel: 0.1640 - val_loss: 1.5427 - val_sparse_categorical_accuracy: 0.5180 - val_std_of_conv2d_2_kernel: 0.1692\n",
            "Epoch 4/5\n",
            "156/156 [==============================] - 14s 91ms/step - loss: 0.1226 - sparse_categorical_accuracy: 0.9641 - std_of_conv2d_2_kernel: 0.1733 - val_loss: 0.5243 - val_sparse_categorical_accuracy: 0.8921 - val_std_of_conv2d_2_kernel: 0.1767\n",
            "Epoch 5/5\n",
            "156/156 [==============================] - 14s 89ms/step - loss: 0.1064 - sparse_categorical_accuracy: 0.9692 - std_of_conv2d_2_kernel: 0.1806 - val_loss: 12.6694 - val_sparse_categorical_accuracy: 0.0982 - val_std_of_conv2d_2_kernel: 0.1842\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tensorflow.python.keras.callbacks.History at 0x7fdc4873ad30>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 22
        }
      ]
    }
  ]
}