tf-levenberg-marquardt.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "tf-levenberg-marquardt.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "FEf5CbTCVa5p"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CMfWCCyHvEjV"
      },
      "source": [
        "# Setup Levenberg-Marquardt\r\n",
        "### Install dependencies from GitHub"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Jzb-I0BbcSkZ"
      },
      "source": [
        "import os\r\n",
        "\r\n",
        "os.chdir(\"/content\")\r\n",
        "\r\n",
        "lm_dir = \"tf-levenberg-marquardt\"\r\n",
        "if not os.path.exists(lm_dir):\r\n",
        "  !git clone https://github.com/fabiodimarco/$lm_dir\r\n",
        "\r\n",
        "os.chdir(lm_dir)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "unTfqvN2VPoB"
      },
      "source": [
        "# Test curve fitting\n",
        "\n",
        "The function `y = sinc(10 * x)` is fitted using a Shallow Neural Network with 61 parameters.\n",
        "Despite the triviality of the problem, first-order methods such as Adam fail to converge, while Levenberg–Marquardt converges rapidly with very low loss values. The values of learning_rate were chosen experimentally on the basis of the results obtained by each algorithm."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZLMXJYbEThWC",
        "cellView": "both"
      },
      "source": [
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import time\n",
        "import levenberg_marquardt as lm\n",
        "\n",
        "input_size = 20000\n",
        "batch_size = 1000\n",
        "\n",
        "x_train = np.linspace(-1, 1, input_size, dtype=np.float64)\n",
        "y_train = np.sinc(10 * x_train)\n",
        "\n",
        "x_train = tf.expand_dims(tf.cast(x_train, tf.float32), axis=-1)\n",
        "y_train = tf.expand_dims(tf.cast(y_train, tf.float32), axis=-1)\n",
        "\n",
        "train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_dataset = train_dataset.shuffle(input_size)\n",
        "train_dataset = train_dataset.batch(batch_size).cache()\n",
        "train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "model = tf.keras.Sequential([\n",
        "    tf.keras.layers.Dense(20, activation='tanh', input_shape=(1,)),\n",
        "    tf.keras.layers.Dense(1, activation='linear')])\n",
        "\n",
        "model.summary()\n",
        "\n",
        "model.compile(\n",
        "    optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),\n",
        "    loss=tf.keras.losses.MeanSquaredError())\n",
        "\n",
        "model_wrapper = lm.ModelWrapper(\n",
        "    tf.keras.models.clone_model(model))\n",
        "\n",
        "model_wrapper.compile(\n",
        "    optimizer=tf.keras.optimizers.SGD(learning_rate=1.0),\n",
        "    loss=lm.MeanSquaredError())\n",
        "\n",
        "print(\"Train using Adam\")\n",
        "t1_start = time.perf_counter()\n",
        "model.fit(train_dataset, epochs=1000)\n",
        "t1_stop = time.perf_counter()\n",
        "print(\"Elapsed time: \", t1_stop - t1_start)\n",
        "\n",
        "print(\"\\n_________________________________________________________________\")\n",
        "print(\"Train using Levenberg-Marquardt\")\n",
        "t2_start = time.perf_counter()\n",
        "model_wrapper.fit(train_dataset, epochs=100)\n",
        "t2_stop = time.perf_counter()\n",
        "print(\"Elapsed time: \", t2_stop - t2_start)\n",
        "\n",
        "print(\"\\n_________________________________________________________________\")\n",
        "print(\"Plot results\")\n",
        "plt.plot(x_train, y_train, 'b-', label=\"reference\")\n",
        "plt.plot(x_train, model.predict(x_train), 'g--', label=\"adam\")\n",
        "plt.plot(x_train, model_wrapper.predict(x_train), 'r--', label=\"lm\")\n",
        "plt.legend()\n",
        "plt.show()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BfEHkJ_VVWdf"
      },
      "source": [
        "# Test mnist classification\n",
        "\n",
        "The classification is performed using a Convolutional Neural Network with 1026 parameters.\n",
        "This time there were no particular benefits as in the previous case. Even if Levenberg–Marquardt converges with far fewer epochs than Adam, the longer execution time per step nullifies its advantages.\n",
        "However, both methods achieve roughly the same accuracy values on train and test set."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RnvOIrjSkcEo",
        "cellView": "both"
      },
      "source": [
        "import tensorflow as tf\n",
        "import time\n",
        "import levenberg_marquardt as lm\n",
        "\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "\n",
        "x_train = tf.cast(x_train / 255.0, dtype=tf.float32)\n",
        "x_test = tf.cast(x_test / 255.0, dtype=tf.float32)\n",
        "\n",
        "x_train = tf.expand_dims(x_train, axis=-1)\n",
        "x_test = tf.expand_dims(x_test, axis=-1)\n",
        "\n",
        "y_train = tf.cast(y_train, dtype=tf.float32)\n",
        "y_test = tf.cast(y_test, dtype=tf.float32)\n",
        "\n",
        "train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_dataset = train_dataset.shuffle(60000)\n",
        "train_dataset = train_dataset.batch(6000).cache()\n",
        "train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "model = tf.keras.models.Sequential([\n",
        "    tf.keras.layers.Conv2D(filters=8, kernel_size=4, strides=2, padding='valid',\n",
        "                           activation='tanh', input_shape=(28, 28, 1)),\n",
        "    tf.keras.layers.Conv2D(filters=4, kernel_size=4, strides=2, padding='valid',\n",
        "                           activation='tanh'),\n",
        "    tf.keras.layers.Conv2D(filters=4, kernel_size=2, strides=1, padding='valid',\n",
        "                           activation='tanh'),\n",
        "    tf.keras.layers.Conv2D(filters=4, kernel_size=2, strides=1, padding='valid',\n",
        "                           activation='tanh'),\n",
        "    tf.keras.layers.Conv2D(filters=4, kernel_size=2, strides=1, padding='valid',\n",
        "                           activation='tanh'),\n",
        "    tf.keras.layers.Flatten(),\n",
        "    tf.keras.layers.Dense(10, activation='linear')\n",
        "])\n",
        "\n",
        "model.summary()\n",
        "\n",
        "model.compile(\n",
        "    optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),\n",
        "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    metrics=[\"accuracy\"])\n",
        "\n",
        "model_wrapper = lm.ModelWrapper(tf.keras.models.clone_model(model))\n",
        "\n",
        "model_wrapper.compile(\n",
        "    optimizer=tf.keras.optimizers.SGD(learning_rate=0.1),\n",
        "    loss=lm.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    solve_method='solve',\n",
        "    metrics=['accuracy'])\n",
        "\n",
        "print(\"Train using Adam\")\n",
        "t1_start = time.perf_counter()\n",
        "model.fit(train_dataset, epochs=200)\n",
        "t1_stop = time.perf_counter()\n",
        "print(\"Elapsed time: \", t1_stop - t1_start)\n",
        "\n",
        "print(\"\\n_________________________________________________________________\")\n",
        "print(\"Train using Levenberg-Marquardt\")\n",
        "t2_start = time.perf_counter()\n",
        "model_wrapper.fit(train_dataset, epochs=10)\n",
        "t2_stop = time.perf_counter()\n",
        "print(\"Elapsed time: \", t2_stop - t2_start)\n",
        "\n",
        "print(\"\\n_________________________________________________________________\")\n",
        "print(\"Test set results\")\n",
        "\n",
        "test_loss, test_acc = model.evaluate(x=x_test, y=y_test, verbose=0)\n",
        "print(\"adam - test_loss: %f - test_accuracy: %f\" % (test_loss, test_acc))\n",
        "\n",
        "test_loss, test_acc = model_wrapper.evaluate(x=x_test, y=y_test, verbose=0)\n",
        "print(\"lm - test_loss: %f - test_accuracy: %f\" % (test_loss, test_acc))\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}