{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Port TensorFlow Quickstart to NNI\nThis is a modified version of `TensorFlow quickstart`_.\n\nIt can be run directly and will have the exact same result as original version.\n\nFurthermore, it enables the ability of auto tuning with an NNI *experiment*, which will be detailed later.\n\nIt is recommended to run this script directly first to verify the environment.\n\nThere are 3 key differences from the original version:\n\n1. In `Get optimized hyperparameters`_ part, it receives generated hyperparameters.\n2. In `(Optional) Report intermediate results`_ part, it reports per-epoch accuracy metrics.\n3. In `Report final result`_ part, it reports final accuracy.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import nni\nimport tensorflow as tf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Hyperparameters to be tuned\nThese are the hyperparameters that will be tuned later.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "params = {\n    'dense_units': 128,\n    'activation_type': 'relu',\n    'dropout_rate': 0.2,\n    'learning_rate': 0.001,\n}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Get optimized hyperparameters\nIf run directly, :func:`nni.get_next_parameter` is a no-op and returns an empty dict.\nBut with an NNI *experiment*, it will receive optimized hyperparameters from tuning algorithm.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "optimized_params = nni.get_next_parameter()\nparams.update(optimized_params)\nprint(params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Load dataset\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "mnist = tf.keras.datasets.mnist\n\n(x_train, y_train), (x_test, y_test) = mnist.load_data()\nx_train, x_test = x_train / 255.0, x_test / 255.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Build model with hyperparameters\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "model = tf.keras.models.Sequential([\n    tf.keras.layers.Flatten(input_shape=(28, 28)),\n    tf.keras.layers.Dense(params['dense_units'], activation=params['activation_type']),\n    tf.keras.layers.Dropout(params['dropout_rate']),\n    tf.keras.layers.Dense(10)\n])\n\nadam = tf.keras.optimizers.Adam(learning_rate=params['learning_rate'])\nloss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\nmodel.compile(optimizer=adam, loss=loss_fn, metrics=['accuracy'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## (Optional) Report intermediate results\nThe callback reports per-epoch accuracy to show learning curve in the web portal.\nYou can also leverage the metrics for early stopping with :doc:`NNI assessors </hpo/assessors>`.\n\nThis part can be safely skipped and the experiment will work fine.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "callback = tf.keras.callbacks.LambdaCallback(\n    on_epoch_end = lambda epoch, logs: nni.report_intermediate_result(logs['accuracy'])\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Train and evluate the model\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "model.fit(x_train, y_train, epochs=5, verbose=2, callbacks=[callback])\nloss, accuracy = model.evaluate(x_test, y_test, verbose=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Report final result\nReport final accuracy to NNI so the tuning algorithm can suggest better hyperparameters.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nni.report_final_result(accuracy)"
      ]
    }
  ],
  "metadata": {
    "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.3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}