tensorflow-handbook-swift-example.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Swift for Tensorflow MNIST Example.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Swift",
      "language": "swift",
      "name": "swift"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/huan/tensorflow-handbook-swift/blob/master/tensorflow-handbook-swift-example.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gQA8-8fHTSY7",
        "colab_type": "text"
      },
      "source": [
        "# Swift MNIST Example\n",
        "\n",
        "Learn more from Github: https://github.com/huan/tensorflow-handbook-swift\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gF303ze8dxXv",
        "colab_type": "text"
      },
      "source": [
        "## MNIST Dataset\n",
        "\n",
        "<https://github.com/huan/swift-MNIST/>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m0Ce4qO38-pU",
        "colab_type": "code",
        "outputId": "3d744f31-b8c5-4e73-8414-f1d73d6e64a5",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 272
        }
      },
      "source": [
        "%install-location $cwd/swift-install\n",
        "%install '.package(url: \"https://github.com/huan/swift-MNIST\", from: \"0.4.0\")' MNIST"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Installing packages:\n",
            "\t.package(url: \"https://github.com/huan/swift-MNIST\", from: \"0.4.0\")\n",
            "\t\tMNIST\n",
            "With SwiftPM flags: []\n",
            "Working in: /tmp/tmpdu69j2gk/swift-install\n",
            "Fetching https://github.com/huan/swift-MNIST\n",
            "Cloning https://github.com/huan/swift-MNIST\n",
            "Resolving https://github.com/huan/swift-MNIST at 0.4.0\n",
            "[1/2] Compiling MNIST MNIST.swift\n",
            "[2/3] Merging module MNIST\n",
            "[3/4] Compiling jupyterInstalledPackages jupyterInstalledPackages.swift\n",
            "[4/5] Merging module jupyterInstalledPackages\n",
            "[5/5] Linking libjupyterInstalledPackages.so\n",
            "Initializing Swift...\n",
            "Installation complete!\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MBwX4iMZz_dL",
        "colab_type": "text"
      },
      "source": [
        "## Define a Simple MLP Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NODylsfzz7Ls",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import TensorFlow\n",
        "import Python\n",
        "import Foundation\n",
        "\n",
        "struct MLP: Layer {\n",
        "  typealias Input = Tensor<Float>\n",
        "  typealias Output = Tensor<Float>\n",
        "\n",
        "  var flatten = Flatten<Float>()\n",
        "  var dense = Dense<Float>(inputSize: 784, outputSize: 10)\n",
        "  \n",
        "  @differentiable\n",
        "  public func callAsFunction(_ input: Input) -> Output {\n",
        "    var x = input\n",
        "    x = flatten(x)\n",
        "    x = dense(x)\n",
        "    return x\n",
        "  }  \n",
        "}\n",
        "\n",
        "var model = MLP()\n",
        "let optimizer = Adam(for: model)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lV2RzLtMeBBH",
        "colab_type": "text"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GzK9FPV3daq5",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 102
        },
        "outputId": "e2e635e1-786d-4f68-d700-0bccb0bf1dc0"
      },
      "source": [
        "import MNIST\n",
        "\n",
        "let mnist = MNIST()\n",
        "let ((trainImages, trainLabels), (testImages, testLabels)) = mnist.loadData()\n",
        "\n",
        "let imageBatch = Dataset(elements: trainImages).batched(32)\n",
        "let labelBatch = Dataset(elements: trainLabels).batched(32)\n",
        "\n",
        "for (X, y) in zip(imageBatch, labelBatch) {\n",
        "  // Caculate the gradient\n",
        "  // let (_loss, grads) = valueWithGradient(at: model) { model -> Tensor<Float> in\n",
        "  let grads = gradient(at: model) { model -> Tensor<Float> in\n",
        "    let logits = model(X)\n",
        "    return softmaxCrossEntropy(logits: logits, labels: y)\n",
        "  }\n",
        "\n",
        "  // Update parameters by optimizer\n",
        "  optimizer.update(&model.allDifferentiableVariables, along: grads)    \n",
        "}\n",
        "\n",
        "let logits = model(testImages)\n",
        "let acc = mnist.getAccuracy(y: testLabels, logits: logits)\n",
        "\n",
        "print(\"Test Accuracy: \\(acc)\" )"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading train-images-idx3-ubyte ...\n",
            "Downloading train-labels-idx1-ubyte ...\n",
            "Reading data.\n",
            "Constructing data tensors.\n",
            "Test Accuracy: 0.9125\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O5UpckxEsqNO",
        "colab_type": "text"
      },
      "source": [
        "- Credit: This example is inspired from [A set of notebooks explaining swift for tensorflow optimized to run in Google Collaboratory.](https://github.com/zaidalyafeai/Swift4TF)\n",
        "- License [Apache-2.0](https://github.com/tensorflow/swift-models/blob/stable/LICENSE)"
      ]
    }
  ]
}