docs/site/tutorials/introducing_x10.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9TV7IYeqifSv"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors. [Licensed under the Apache License, Version 2.0](#scrollTo=ByZjmtFgB_Y5)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kZRlD4utdPuX"
      },
      "outputs": [],
      "source": [
        "%install '.package(url: \"https://github.com/tensorflow/swift-models\", .branch(\"tensorflow-0.12\"))' Datasets ImageClassificationModels\n",
        "print(\"\\u{001B}[2J\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tRIJp_4m_Afz"
      },
      "outputs": [],
      "source": [
        "// #@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "// Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "// you may not use this file except in compliance with the License.\n",
        "// You may obtain a copy of the License at\n",
        "//\n",
        "// https://www.apache.org/licenses/LICENSE-2.0\n",
        "//\n",
        "// Unless required by applicable law or agreed to in writing, software\n",
        "// distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "// See the License for the specific language governing permissions and\n",
        "// limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sI1ZtrdiA4aY"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        " <td>\n",
        "  <a target=\"_blank\" href=\"https://www.tensorflow.org/swift/tutorials/introducing_x10\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        " </td>\n",
        " <td>\n",
        "  <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/swift/blob/main/docs/site/tutorials/introducing_x10.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        " </td>\n",
        " <td>\n",
        "  <a target=\"_blank\" href=\"https://github.com/tensorflow/swift/blob/main/docs/site/tutorials/introducing_x10.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        " </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8sa42_NblqRE"
      },
      "source": [
        "# Introducing X10\n",
        "\n",
        "By default, Swift For TensorFlow performs tensor operations using eager dispatch. This allows for rapid iteration, but isn't the most performant option for training machine learning models.\n",
        "\n",
        "The [X10 tensor library](https://github.com/tensorflow/swift-apis/blob/main/Documentation/X10/API_GUIDE.md) adds a high-performance backend to Swift for TensorFlow, leveraging tensor tracing and the [XLA compiler](https://www.tensorflow.org/xla). This tutorial will introduce X10 and guide you through the process of updating a training loop to run on GPUs or TPUs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W7MpNcIwIIy8"
      },
      "source": [
        "## Eager vs. X10 tensors"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lM9dRji7IIy8"
      },
      "source": [
        "Accelerated calculations in Swift for TensorFlow are performed through the Tensor type. Tensors can participate in a wide variety of operations, and are the fundamental building blocks of machine learning models.\n",
        "\n",
        "By default, a Tensor uses eager execution to perform calculations on an operation-by-operation basis. Each Tensor has an associated Device that describes what hardware it is attached to and what backend is used for it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OHRTNQJo1TxT"
      },
      "outputs": [],
      "source": [
        "import TensorFlow\n",
        "import Foundation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FCMWR11NIIy-"
      },
      "outputs": [],
      "source": [
        "let eagerTensor1 = Tensor([0.0, 1.0, 2.0])\n",
        "let eagerTensor2 = Tensor([1.5, 2.5, 3.5])\n",
        "let eagerTensorSum = eagerTensor1 + eagerTensor2\n",
        "print(eagerTensorSum)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1qad9_yMYf6F"
      },
      "outputs": [],
      "source": [
        "print(eagerTensor1.device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HrlMNOinIIy_"
      },
      "source": [
        "If you are running this notebook on a GPU-enabled instance, you should see that hardware reflected in the device description above. The eager runtime does not have support for TPUs, so if you are using one of them as an accelerator you will see the CPU being used as a hardware target."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eoyLeSQVIIy9"
      },
      "source": [
        "When creating a Tensor, the default eager mode device can be overridden by specifying an alternative. This is how you opt-in to performing calculations using the X10 backend."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NrRQhQaHaJm9"
      },
      "outputs": [],
      "source": [
        "let x10Tensor1 = Tensor([0.0, 1.0, 2.0], on: Device.defaultXLA)\n",
        "let x10Tensor2 = Tensor([1.5, 2.5, 3.5], on: Device.defaultXLA)\n",
        "let x10TensorSum = x10Tensor1 + x10Tensor2\n",
        "print(x10TensorSum)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VbqeudQCaqwv"
      },
      "outputs": [],
      "source": [
        "print(x10Tensor1.device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mIbIOW0HIIzA"
      },
      "source": [
        "If you're running this in a GPU-enabled instance, you should see that accelerator listed in the X10 tensor's device. Unlike for eager execution, if you are running this in a TPU-enabled instance, you should now see that calculations are using that device. X10 is how you take advantage of TPUs within Swift for TensorFlow.\n",
        "\n",
        "The default eager and X10 devices will attempt to use the first accelerator on the system. If you have GPUs attached, the will use the first available GPU. If TPUs are present, X10 will use the first TPU core by default. If no accelerator is found or supported, the default device will fall back to the CPU.\n",
        "\n",
        "Beyond the default eager and XLA devices, you can provide specific hardware and backend targets in a Device:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "De59VwJ35SvG"
      },
      "outputs": [],
      "source": [
        "// let tpu1 = Device(kind: .TPU, ordinal: 1, backend: .XLA)\n",
        "// let tpuTensor1 = Tensor([0.0, 1.0, 2.0], on: tpu1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rU0WY_sJodio"
      },
      "source": [
        "## Training an eager-mode model\n",
        "\n",
        "Let's take a look at how you'd set up and train a model using the default eager execution mode. In this example, we'll be using the simple LeNet-5 model from the [swift-models repository](https://github.com/tensorflow/swift-models) and the MNIST handwritten digit classification dataset.\n",
        "\n",
        "First, we'll set up and download the MNIST dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kqXILiXhq-iM"
      },
      "outputs": [],
      "source": [
        "import Datasets\n",
        "\n",
        "let epochCount = 5\n",
        "let batchSize = 128\n",
        "let dataset = MNIST(batchSize: batchSize)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1pMewsl0VgnJ"
      },
      "source": [
        "Next, we will configure the model and optimizer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fEAEyUExXT3I"
      },
      "outputs": [],
      "source": [
        "import ImageClassificationModels\n",
        "\n",
        "var eagerModel = LeNet()\n",
        "var eagerOptimizer = SGD(for: eagerModel, learningRate: 0.1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cwNDwzS2QgS1"
      },
      "source": [
        "Now, we will implement basic progress tracking and reporting.  All intermediate statistics are kept as tensors on the same device where training is run and `scalarized()` is called only during reporting. This will be especially important later when using X10, because it avoids unnecessary materialization of lazy tensors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eNzOdly3QY_P"
      },
      "outputs": [],
      "source": [
        "struct Statistics {\n",
        "    var correctGuessCount = Tensor<Int32>(0, on: Device.default)\n",
        "    var totalGuessCount = Tensor<Int32>(0, on: Device.default)\n",
        "    var totalLoss = Tensor<Float>(0, on: Device.default)\n",
        "    var batches: Int = 0\n",
        "    var accuracy: Float { \n",
        "        Float(correctGuessCount.scalarized()) / Float(totalGuessCount.scalarized()) * 100 \n",
        "    } \n",
        "    var averageLoss: Float { totalLoss.scalarized() / Float(batches) }\n",
        "\n",
        "    init(on device: Device = Device.default) {\n",
        "        correctGuessCount = Tensor<Int32>(0, on: device)\n",
        "        totalGuessCount = Tensor<Int32>(0, on: device)\n",
        "        totalLoss = Tensor<Float>(0, on: device)\n",
        "    }\n",
        "\n",
        "    mutating func update(logits: Tensor<Float>, labels: Tensor<Int32>, loss: Tensor<Float>) {\n",
        "        let correct = logits.argmax(squeezingAxis: 1) .== labels\n",
        "        correctGuessCount += Tensor<Int32>(correct).sum()\n",
        "        totalGuessCount += Int32(labels.shape[0])\n",
        "        totalLoss += loss\n",
        "        batches += 1\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w3lmTRCWT5sS"
      },
      "source": [
        "Finally, we'll run the model through a training loop for five epochs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W9bUsiOxVf_v"
      },
      "outputs": [],
      "source": [
        "print(\"Beginning training...\")\n",
        "\n",
        "for (epoch, batches) in dataset.training.prefix(epochCount).enumerated() {\n",
        "    let start = Date()\n",
        "    var trainStats = Statistics()\n",
        "    var testStats = Statistics()\n",
        "    \n",
        "    Context.local.learningPhase = .training\n",
        "    for batch in batches {\n",
        "        let (images, labels) = (batch.data, batch.label)\n",
        "        let 𝛁model = TensorFlow.gradient(at: eagerModel) { eagerModel -> Tensor<Float> in\n",
        "            let ŷ = eagerModel(images)\n",
        "            let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)\n",
        "            trainStats.update(logits: ŷ, labels: labels, loss: loss)\n",
        "            return loss\n",
        "        }\n",
        "        eagerOptimizer.update(&eagerModel, along: 𝛁model)\n",
        "    }\n",
        "\n",
        "    Context.local.learningPhase = .inference\n",
        "    for batch in dataset.validation {\n",
        "        let (images, labels) = (batch.data, batch.label)\n",
        "        let ŷ = eagerModel(images)\n",
        "        let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)\n",
        "        testStats.update(logits: ŷ, labels: labels, loss: loss)\n",
        "    }\n",
        "\n",
        "    print(\n",
        "        \"\"\"\n",
        "        [Epoch \\(epoch)] \\\n",
        "        Training Loss: \\(String(format: \"%.3f\", trainStats.averageLoss)), \\\n",
        "        Training Accuracy: \\(trainStats.correctGuessCount)/\\(trainStats.totalGuessCount) \\\n",
        "        (\\(String(format: \"%.1f\", trainStats.accuracy))%), \\\n",
        "        Test Loss: \\(String(format: \"%.3f\", testStats.averageLoss)), \\\n",
        "        Test Accuracy: \\(testStats.correctGuessCount)/\\(testStats.totalGuessCount) \\\n",
        "        (\\(String(format: \"%.1f\", testStats.accuracy))%) \\\n",
        "        seconds per epoch: \\(String(format: \"%.1f\", Date().timeIntervalSince(start)))\n",
        "        \"\"\")\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ED0HGZW2gWY"
      },
      "source": [
        "As you can see, the model trained as we would expect, and its accuracy against the validation set increased each epoch. This is how Swift for TensorFlow models are defined and run using eager execution, now let's see what modifications need to be made to take advantage of X10."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Te7sNNx9c_am"
      },
      "source": [
        "## Training an X10 model\n",
        "\n",
        "Datasets, models, and optimizers contain tensors that are initialized on the default eager execution device. To work with X10, we'll need to move these tensors to an X10 device."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MaN7fM-lAe7r"
      },
      "outputs": [],
      "source": [
        "let device = Device.defaultXLA\n",
        "print(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cJfSg0wDAgC7"
      },
      "source": [
        "For the datasets, we'll do that at the point in which batches are processed in the training loop, so we can re-use the dataset from the eager execution model.\n",
        "\n",
        "In the case of the model and optimizer, we'll initialize them with their internal tensors on the eager execution device, then move them over to the X10 device."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jpcOByipc75O"
      },
      "outputs": [],
      "source": [
        "var x10Model = LeNet()\n",
        "x10Model.move(to: device)\n",
        "\n",
        "var x10Optimizer = SGD(for: x10Model, learningRate: 0.1)\n",
        "x10Optimizer = SGD(copying: x10Optimizer, to: device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hQvza3dUXlr0"
      },
      "source": [
        "The modifications needed for the training loop come at a few specific points. First, we'll need to move the batches of training data over to the X10 device. This is done via `Tensor(copying:to:)` when each batch is retrieved.\n",
        "\n",
        "The next change is to indicate where to cut off the traces during the training loop. X10 works by tracing through the tensor calculations needed in your code and just-in-time compiling an optimized representation of that trace. In the case of a training loop, you’re repeating the same operation over and over again, an ideal section to trace, compile, and re-use.\n",
        "\n",
        "In the absence of code that explicitly requests a value from a Tensor (these usually stand out as `.scalars` or `.scalarized()` calls), X10 will attempt to compile all loop iterations together. To prevent this, and cut the trace at a specific point, we place an explicit `LazyTensorBarrier()` after the optimizer updates the model weights and after the loss and accuracy are obtained during validation. This creates two reused traces: each step in the training loop and each batch of inference during validation.\n",
        "\n",
        "These modifications result in the following training loop."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XrZee8n3Y17_"
      },
      "outputs": [],
      "source": [
        "print(\"Beginning training...\")\n",
        "\n",
        "for (epoch, batches) in dataset.training.prefix(epochCount).enumerated() {\n",
        "    let start = Date()\n",
        "    var trainStats = Statistics(on: device)\n",
        "    var testStats = Statistics(on: device)\n",
        "    \n",
        "    Context.local.learningPhase = .training\n",
        "    for batch in batches {\n",
        "        let (eagerImages, eagerLabels) = (batch.data, batch.label)\n",
        "        let images = Tensor(copying: eagerImages, to: device)\n",
        "        let labels = Tensor(copying: eagerLabels, to: device)\n",
        "        let 𝛁model = TensorFlow.gradient(at: x10Model) { x10Model -> Tensor<Float> in\n",
        "            let ŷ = x10Model(images)\n",
        "            let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)\n",
        "            trainStats.update(logits: ŷ, labels: labels, loss: loss)\n",
        "            return loss\n",
        "        }\n",
        "        x10Optimizer.update(&x10Model, along: 𝛁model)\n",
        "        LazyTensorBarrier()\n",
        "    }\n",
        "\n",
        "    Context.local.learningPhase = .inference\n",
        "    for batch in dataset.validation {\n",
        "        let (eagerImages, eagerLabels) = (batch.data, batch.label)\n",
        "        let images = Tensor(copying: eagerImages, to: device)\n",
        "        let labels = Tensor(copying: eagerLabels, to: device)\n",
        "        let ŷ = x10Model(images)\n",
        "        let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)\n",
        "        LazyTensorBarrier()\n",
        "        testStats.update(logits: ŷ, labels: labels, loss: loss)\n",
        "    }\n",
        "\n",
        "    print(\n",
        "        \"\"\"\n",
        "        [Epoch \\(epoch)] \\\n",
        "        Training Loss: \\(String(format: \"%.3f\", trainStats.averageLoss)), \\\n",
        "        Training Accuracy: \\(trainStats.correctGuessCount)/\\(trainStats.totalGuessCount) \\\n",
        "        (\\(String(format: \"%.1f\", trainStats.accuracy))%), \\\n",
        "        Test Loss: \\(String(format: \"%.3f\", testStats.averageLoss)), \\\n",
        "        Test Accuracy: \\(testStats.correctGuessCount)/\\(testStats.totalGuessCount) \\\n",
        "        (\\(String(format: \"%.1f\", testStats.accuracy))%) \\\n",
        "        seconds per epoch: \\(String(format: \"%.1f\", Date().timeIntervalSince(start)))\n",
        "        \"\"\")\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qej_Z6V3mZnG"
      },
      "source": [
        "Training of the model using the X10 backend should have proceeded in the same manner as the eager execution model did before. You may have noticed a delay before the first batch and at the end of the first epoch, due to the just-in-time compilation of the unique traces at those points. If you're running this with an accelerator attached, you should have seen the training after that point proceeding faster than with eager mode.\n",
        "\n",
        "There is a tradeoff of initial trace compilation time vs. faster throughput, but in most machine learning models the increase in throughput from repeated operations should more than offset compilation overhead. In practice, we've seen an over 4X improvement in throughput with X10 in some training cases.\n",
        "\n",
        "As has been stated before, using X10 now makes it not only possible but easy to work with TPUs, unlocking that whole class of accelerators for your Swift for TensorFlow models."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "introducing_x10.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Swift",
      "name": "swift"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}