docs/site/tutorials/model_training_walkthrough.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QyCcF45zBQ3E"
   },
   "source": [
    "##### Copyright 2018 - 2020 The TensorFlow Authors. [Licensed under the Apache License, Version 2.0](#scrollTo=y_UVSRtBBsJk)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "CPII1rGR2rF9"
   },
   "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": "zBH72IXMJ3JJ"
   },
   "source": [
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://www.tensorflow.org/swift/tutorials/model_training_walkthrough\"><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/model_training_walkthrough.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/model_training_walkthrough.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": "JtEZ1pCPn--z"
   },
   "source": [
    "# Model training walkthrough"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LDrzLFXE8T1l"
   },
   "source": [
    "This guide introduces Swift for TensorFlow by building a machine learning model that categorizes iris flowers by species. It uses Swift for TensorFlow to:\n",
    "1. Build a model,\n",
    "2. Train this model on example data, and\n",
    "3. Use the model to make predictions about unknown data.\n",
    "\n",
    "## TensorFlow programming\n",
    "\n",
    "This guide uses these high-level Swift for TensorFlow concepts:\n",
    "\n",
    "* Import data with the Epochs API.\n",
    "* Build models using Swift abstractions.\n",
    "* Use Python libraries using Swift's Python interoperability when pure Swift libraries are not available.\n",
    "\n",
    "This tutorial is structured like many TensorFlow programs:\n",
    "\n",
    "1. Import and parse the data sets.\n",
    "2. Select the type of model.\n",
    "3. Train the model.\n",
    "4. Evaluate the model's effectiveness.\n",
    "5. Use the trained model to make predictions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yNr7H-AIoLOR"
   },
   "source": [
    "## Setup program"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1J3AuPBT9gyR"
   },
   "source": [
    "### Configure imports\n",
    "\n",
    "Import TensorFlow and some useful Python modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "g4Wzg69bnwK2"
   },
   "outputs": [],
   "source": [
    "import TensorFlow\n",
    "import PythonKit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5ms1o2W-DF1g"
   },
   "outputs": [],
   "source": [
    "// This cell is here to display the plots in a Jupyter Notebook.\n",
    "// Do not copy it into another environment.\n",
    "%include \"EnableIPythonDisplay.swift\"\n",
    "print(IPythonDisplay.shell.enable_matplotlib(\"inline\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "82TJnHsCY02t"
   },
   "outputs": [],
   "source": [
    "let plt = Python.import(\"matplotlib.pyplot\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-5p_tYTyz43Q"
   },
   "outputs": [],
   "source": [
    "import Foundation\n",
    "import FoundationNetworking\n",
    "func download(from sourceString: String, to destinationString: String) {\n",
    "    let source = URL(string: sourceString)!\n",
    "    let destination = URL(fileURLWithPath: destinationString)\n",
    "    let data = try! Data.init(contentsOf: source)\n",
    "    try! data.write(to: destination)\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Zx7wc0LuuxaJ"
   },
   "source": [
    "## The iris classification problem\n",
    "\n",
    "Imagine you are a botanist seeking an automated way to categorize each iris flower you find. Machine learning provides many algorithms to classify flowers statistically. For instance, a sophisticated machine learning program could classify flowers based on photographs. Our ambitions are more modest—we're going to classify iris flowers based on the length and width measurements of their [sepals](https://en.wikipedia.org/wiki/Sepal) and [petals](https://en.wikipedia.org/wiki/Petal).\n",
    "\n",
    "The Iris genus entails about 300 species, but our program will only classify the following three:\n",
    "\n",
    "* Iris setosa\n",
    "* Iris virginica\n",
    "* Iris versicolor\n",
    "\n",
    "<table>\n",
    "  <tr><td>\n",
    "    <img src=\"https://www.tensorflow.org/images/iris_three_species.jpg\"\n",
    "         alt=\"Petal geometry compared for three iris species: Iris setosa, Iris virginica, and Iris versicolor\">\n",
    "  </td></tr>\n",
    "  <tr><td align=\"center\">\n",
    "    <b>Figure 1.</b> <a href=\"https://commons.wikimedia.org/w/index.php?curid=170298\">Iris setosa</a> (by <a href=\"https://commons.wikimedia.org/wiki/User:Radomil\">Radomil</a>, CC BY-SA 3.0), <a href=\"https://commons.wikimedia.org/w/index.php?curid=248095\">Iris versicolor</a>, (by <a href=\"https://commons.wikimedia.org/wiki/User:Dlanglois\">Dlanglois</a>, CC BY-SA 3.0), and <a href=\"https://www.flickr.com/photos/33397993@N05/3352169862\">Iris virginica</a> (by <a href=\"https://www.flickr.com/photos/33397993@N05\">Frank Mayfield</a>, CC BY-SA 2.0).<br/>&nbsp;\n",
    "  </td></tr>\n",
    "</table>\n",
    "\n",
    "Fortunately, someone has already created a [data set of 120 iris flowers](https://en.wikipedia.org/wiki/Iris_flower_data_set) with the sepal and petal measurements. This is a classic dataset that is popular for beginner machine learning classification problems."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3Px6KAg0Jowz"
   },
   "source": [
    "## Import and parse the training dataset\n",
    "\n",
    "Download the dataset file and convert it into a structure that can be used by this Swift program.\n",
    "\n",
    "### Download the dataset\n",
    "\n",
    "Download the training dataset file from http://download.tensorflow.org/data/iris_training.csv."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DKkgac4WO0mP"
   },
   "outputs": [],
   "source": [
    "let trainDataFilename = \"iris_training.csv\"\n",
    "download(from: \"http://download.tensorflow.org/data/iris_training.csv\", to: trainDataFilename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qnX1-aLors4S"
   },
   "source": [
    "### Inspect the data\n",
    "\n",
    "This dataset, `iris_training.csv`, is a plain text file that stores tabular data formatted as comma-separated values (CSV). Let's look a the first 5 entries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FQvb_JYdrpPm"
   },
   "outputs": [],
   "source": [
    "let f = Python.open(trainDataFilename)\n",
    "for _ in 0..<5 {\n",
    "    print(Python.next(f).strip())\n",
    "}\n",
    "print(f.close())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kQhzD6P-uBoq"
   },
   "source": [
    "From this view of the dataset, notice the following:\n",
    "\n",
    "1. The first line is a header containing information about the dataset:\n",
    "  * There are 120 total examples. Each example has four features and one of three possible label names. \n",
    "2. Subsequent rows are data records, one *[example](https://developers.google.com/machine-learning/glossary/#example)* per line, where:\n",
    "  * The first four fields are *[features](https://developers.google.com/machine-learning/glossary/#feature)*: these are characteristics of an example. Here, the fields hold float numbers representing flower measurements.\n",
    "  * The last column is the *[label](https://developers.google.com/machine-learning/glossary/#label)*: this is the value we want to predict. For this dataset, it's an integer value of 0, 1, or 2 that corresponds to a flower name.\n",
    "\n",
    "Let's write that out in code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9Edhevw7exl6"
   },
   "outputs": [],
   "source": [
    "let featureNames = [\"sepal_length\", \"sepal_width\", \"petal_length\", \"petal_width\"]\n",
    "let labelName = \"species\"\n",
    "let columnNames = featureNames + [labelName]\n",
    "\n",
    "print(\"Features: \\(featureNames)\")\n",
    "print(\"Label: \\(labelName)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CCtwLoJhhDNc"
   },
   "source": [
    "Each label is associated with string name (for example, \"setosa\"), but machine learning typically relies on numeric values. The label numbers are mapped to a named representation, such as:\n",
    "\n",
    "* `0`: Iris setosa\n",
    "* `1`: Iris versicolor\n",
    "* `2`: Iris virginica\n",
    "\n",
    "For more information about features and labels, see the [ML Terminology section of the Machine Learning Crash Course](https://developers.google.com/machine-learning/crash-course/framing/ml-terminology)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sVNlJlUOhkoX"
   },
   "outputs": [],
   "source": [
    "let classNames = [\"Iris setosa\", \"Iris versicolor\", \"Iris virginica\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dqPkQExM2Pwt"
   },
   "source": [
    "### Create a dataset using the Epochs API\n",
    "\n",
    "Swift for TensorFlow's Epochs API is a high-level API for reading data and transforming it into a form used for training. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bBx_C6UWO0mc"
   },
   "outputs": [],
   "source": [
    "let batchSize = 32\n",
    "\n",
    "/// A batch of examples from the iris dataset.\n",
    "struct IrisBatch {\n",
    "    /// [batchSize, featureCount] tensor of features.\n",
    "    let features: Tensor<Float>\n",
    "\n",
    "    /// [batchSize] tensor of labels.\n",
    "    let labels: Tensor<Int32>\n",
    "}\n",
    "\n",
    "/// Conform `IrisBatch` to `Collatable` so that we can load it into a `TrainingEpoch`.\n",
    "extension IrisBatch: Collatable {\n",
    "    public init<BatchSamples: Collection>(collating samples: BatchSamples)\n",
    "        where BatchSamples.Element == Self {\n",
    "        /// `IrisBatch`es are collated by stacking their feature and label tensors\n",
    "        /// along the batch axis to produce a single feature and label tensor\n",
    "        features = Tensor<Float>(stacking: samples.map{$0.features})\n",
    "        labels = Tensor<Int32>(stacking: samples.map{$0.labels})\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SO6elT-kwwIK"
   },
   "source": [
    "Since the datasets we downloaded are in CSV format, let's write a function to load in the data as a list of IrisBatch objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LwA21wYCguO5"
   },
   "outputs": [],
   "source": [
    "/// Initialize an `IrisBatch` dataset from a CSV file.\n",
    "func loadIrisDatasetFromCSV(\n",
    "        contentsOf: String, hasHeader: Bool, featureColumns: [Int], labelColumns: [Int]) -> [IrisBatch] {\n",
    "        let np = Python.import(\"numpy\")\n",
    "\n",
    "        let featuresNp = np.loadtxt(\n",
    "            contentsOf,\n",
    "            delimiter: \",\",\n",
    "            skiprows: hasHeader ? 1 : 0,\n",
    "            usecols: featureColumns,\n",
    "            dtype: Float.numpyScalarTypes.first!)\n",
    "        guard let featuresTensor = Tensor<Float>(numpy: featuresNp) else {\n",
    "            // This should never happen, because we construct featuresNp in such a\n",
    "            // way that it should be convertible to tensor.\n",
    "            fatalError(\"np.loadtxt result can't be converted to Tensor\")\n",
    "        }\n",
    "\n",
    "        let labelsNp = np.loadtxt(\n",
    "            contentsOf,\n",
    "            delimiter: \",\",\n",
    "            skiprows: hasHeader ? 1 : 0,\n",
    "            usecols: labelColumns,\n",
    "            dtype: Int32.numpyScalarTypes.first!)\n",
    "        guard let labelsTensor = Tensor<Int32>(numpy: labelsNp) else {\n",
    "            // This should never happen, because we construct labelsNp in such a\n",
    "            // way that it should be convertible to tensor.\n",
    "            fatalError(\"np.loadtxt result can't be converted to Tensor\")\n",
    "        }\n",
    "\n",
    "        return zip(featuresTensor.unstacked(), labelsTensor.unstacked()).map{IrisBatch(features: $0.0, labels: $0.1)}\n",
    "\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HbmOmUWneVGM"
   },
   "source": [
    "We can now use the CSV loading function to load the training dataset and create a `TrainingEpochs` object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zFnMejfFYgSV"
   },
   "outputs": [],
   "source": [
    "let trainingDataset: [IrisBatch] = loadIrisDatasetFromCSV(contentsOf: trainDataFilename, \n",
    "                                                  hasHeader: true, \n",
    "                                                  featureColumns: [0, 1, 2, 3], \n",
    "                                                  labelColumns: [4])\n",
    "\n",
    "let trainingEpochs: TrainingEpochs = TrainingEpochs(samples: trainingDataset, batchSize: batchSize)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gB_RSn62c-3G"
   },
   "source": [
    "The `TrainingEpochs` object is an infinite sequence of epochs. Each epoch contains `IrisBatch`es. Let's look at the first element of the first epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iDuG94H-C122"
   },
   "outputs": [],
   "source": [
    "let firstTrainEpoch = trainingEpochs.next()!\n",
    "let firstTrainBatch = firstTrainEpoch.first!.collated\n",
    "let firstTrainFeatures = firstTrainBatch.features\n",
    "let firstTrainLabels = firstTrainBatch.labels\n",
    "\n",
    "print(\"First batch of features: \\(firstTrainFeatures)\")\n",
    "print(\"firstTrainFeatures.shape: \\(firstTrainFeatures.shape)\")\n",
    "print(\"First batch of labels: \\(firstTrainLabels)\")\n",
    "print(\"firstTrainLabels.shape: \\(firstTrainLabels.shape)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "E63mArnQaAGz"
   },
   "source": [
    "Notice that the features for the first `batchSize` examples are grouped together (or *batched*) into `firstTrainFeatures`, and that the labels for the first `batchSize` examples are batched into `firstTrainLabels`.\n",
    "\n",
    "You can start to see some clusters by plotting a few features from the batch, using Python's matplotlib:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "me5Wn-9FcyyO"
   },
   "outputs": [],
   "source": [
    "let firstTrainFeaturesTransposed = firstTrainFeatures.transposed()\n",
    "let petalLengths = firstTrainFeaturesTransposed[2].scalars\n",
    "let sepalLengths = firstTrainFeaturesTransposed[0].scalars\n",
    "\n",
    "plt.scatter(petalLengths, sepalLengths, c: firstTrainLabels.array.scalars)\n",
    "plt.xlabel(\"Petal length\")\n",
    "plt.ylabel(\"Sepal length\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LsaVrtNM3Tx5"
   },
   "source": [
    "## Select the type of model\n",
    "\n",
    "### Why model?\n",
    "\n",
    "A *[model](https://developers.google.com/machine-learning/crash-course/glossary#model)* is a relationship between features and the label.  For the iris classification problem, the model defines the relationship between the sepal and petal measurements and the predicted iris species. Some simple models can be described with a few lines of algebra, but complex machine learning models have a large number of parameters that are difficult to summarize.\n",
    "\n",
    "Could you determine the relationship between the four features and the iris species *without* using machine learning?  That is, could you use traditional programming techniques (for example, a lot of conditional statements) to create a model?  Perhaps—if you analyzed the dataset long enough to determine the relationships between petal and sepal measurements to a particular species. And this becomes difficult—maybe impossible—on more complicated datasets. A good machine learning approach *determines the model for you*. If you feed enough representative examples into the right machine learning model type, the program will figure out the relationships for you.\n",
    "\n",
    "### Select the model\n",
    "\n",
    "We need to select the kind of model to train. There are many types of models and picking a good one takes experience. This tutorial uses a neural network to solve the iris classification problem. *[Neural networks](https://developers.google.com/machine-learning/glossary/#neural_network)* can find complex relationships between features and the label. It is a highly-structured graph, organized into one or more *[hidden layers](https://developers.google.com/machine-learning/glossary/#hidden_layer)*. Each hidden layer consists of one or more *[neurons](https://developers.google.com/machine-learning/glossary/#neuron)*. There are several categories of neural networks and this program uses a dense, or *[fully-connected neural network](https://developers.google.com/machine-learning/glossary/#fully_connected_layer)*: the neurons in one layer receive input connections from *every* neuron in the previous layer. For example, Figure 2 illustrates a dense neural network consisting of an input layer, two hidden layers, and an output layer:\n",
    "\n",
    "<table>\n",
    "  <tr><td>\n",
    "    <img src=\"https://www.tensorflow.org/images/custom_estimators/full_network.png\"\n",
    "         alt=\"A diagram of the network architecture: Inputs, 2 hidden layers, and outputs\">\n",
    "  </td></tr>\n",
    "  <tr><td align=\"center\">\n",
    "    <b>Figure 2.</b> A neural network with features, hidden layers, and predictions.<br/>&nbsp;\n",
    "  </td></tr>\n",
    "</table>\n",
    "\n",
    "When the model from Figure 2 is trained and fed an unlabeled example, it yields three predictions: the likelihood that this flower is the given iris species. This prediction is called *[inference](https://developers.google.com/machine-learning/crash-course/glossary#inference)*. For this example, the sum of the output predictions is 1.0. In Figure 2, this prediction breaks down as: `0.02` for *Iris setosa*, `0.95` for *Iris versicolor*, and `0.03` for *Iris virginica*. This means that the model predicts—with 95% probability—that an unlabeled example flower is an *Iris versicolor*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W23DIMVPQEBt"
   },
   "source": [
    "### Create a model using the Swift for TensorFlow Deep Learning Library\n",
    "\n",
    "The [Swift for TensorFlow Deep Learning Library](https://github.com/tensorflow/swift-apis) defines primitive layers and conventions for wiring them together, which makes it easy to build models and experiment.\n",
    "\n",
    "A model is a `struct` that conforms to [`Layer`](https://www.tensorflow.org/swift/api_docs/Protocols/Layer), which means that it defines a [`callAsFunction(_:)`](https://www.tensorflow.org/swift/api_docs/Protocols/Layer#callasfunction_:) method that maps input `Tensor`s to output `Tensor`s. The `callAsFunction(_:)` method often simply sequences the input through sublayers. Let's define an `IrisModel` that sequences the input through three [`Dense`](https://www.tensorflow.org/swift/api_docs/Structs/Dense) sublayers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wr5A5WvthvZ0"
   },
   "outputs": [],
   "source": [
    "import TensorFlow\n",
    "\n",
    "let hiddenSize: Int = 10\n",
    "struct IrisModel: Layer {\n",
    "    var layer1 = Dense<Float>(inputSize: 4, outputSize: hiddenSize, activation: relu)\n",
    "    var layer2 = Dense<Float>(inputSize: hiddenSize, outputSize: hiddenSize, activation: relu)\n",
    "    var layer3 = Dense<Float>(inputSize: hiddenSize, outputSize: 3)\n",
    "    \n",
    "    @differentiable\n",
    "    func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {\n",
    "        return input.sequenced(through: layer1, layer2, layer3)\n",
    "    }\n",
    "}\n",
    "\n",
    "var model = IrisModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fK0vrIRv_tcc"
   },
   "source": [
    "The activation function determines the output shape of each node in the layer. These non-linearities are important—without them the model would be equivalent to a single layer. There are many available activations, but [ReLU](https://www.tensorflow.org/swift/api_docs/Functions#relu_:) is common for hidden layers.\n",
    "\n",
    "The ideal number of hidden layers and neurons depends on the problem and the dataset. Like many aspects of machine learning, picking the best shape of the neural network requires a mixture of knowledge and experimentation. As a rule of thumb, increasing the number of hidden layers and neurons typically creates a more powerful model, which requires more data to train effectively."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2wFKnhWCpDSS"
   },
   "source": [
    "### Using the model\n",
    "\n",
    "Let's have a quick look at what this model does to a batch of features:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sKjJGIYzO0mr"
   },
   "outputs": [],
   "source": [
    "// Apply the model to a batch of features.\n",
    "let firstTrainPredictions = model(firstTrainFeatures)\n",
    "print(firstTrainPredictions[0..<5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wxyXOhwVr5S3"
   },
   "source": [
    "Here, each example returns a [logit](https://developers.google.com/machine-learning/crash-course/glossary#logits) for each class. \n",
    "\n",
    "To convert these logits to a probability for each class, use the [softmax](https://developers.google.com/machine-learning/crash-course/glossary#softmax) function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_tRwHZmTNTX2"
   },
   "outputs": [],
   "source": [
    "print(softmax(firstTrainPredictions[0..<5]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uRZmchElo481"
   },
   "source": [
    "Taking the `argmax` across classes gives us the predicted class index. But, the model hasn't been trained yet, so these aren't good predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-Jzm_GoErz8B"
   },
   "outputs": [],
   "source": [
    "print(\"Prediction: \\(firstTrainPredictions.argmax(squeezingAxis: 1))\")\n",
    "print(\"    Labels: \\(firstTrainLabels)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vzq2E5J2QMtw"
   },
   "source": [
    "## Train the model\n",
    "\n",
    "*[Training](https://developers.google.com/machine-learning/crash-course/glossary#training)* is the stage of machine learning when the model is gradually optimized, or the model *learns* the dataset. The goal is to learn enough about the structure of the training dataset to make predictions about unseen data. If you learn *too much* about the training dataset, then the predictions only work for the data it has seen and will not be generalizable. This problem is called *[overfitting](https://developers.google.com/machine-learning/crash-course/glossary#overfitting)*—it's like memorizing the answers instead of understanding how to solve a problem.\n",
    "\n",
    "The iris classification problem is an example of *[supervised machine learning](https://developers.google.com/machine-learning/glossary/#supervised_machine_learning)*: the model is trained from examples that contain labels. In *[unsupervised machine learning](https://developers.google.com/machine-learning/glossary/#unsupervised_machine_learning)*, the examples don't contain labels. Instead, the model typically finds patterns among the features."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RaKp8aEjKX6B"
   },
   "source": [
    "### Choose a loss function\n",
    "\n",
    "Both training and evaluation stages need to calculate the model's *[loss](https://developers.google.com/machine-learning/crash-course/glossary#loss)*. This measures how off a model's predictions are from the desired label, in other words, how bad the model is performing. We want to minimize, or optimize, this value.\n",
    "\n",
    "Our model will calculate its loss using the [`softmaxCrossEntropy(logits:labels:)`](https://www.tensorflow.org/swift/api_docs/Functions#/s:10TensorFlow19softmaxCrossEntropy6logits6labelsAA0A0VyxGAG_AFys5Int32VGtAA0aB13FloatingPointRzlF) function which takes the model's class probability predictions and the desired label, and returns the average loss across the examples.\n",
    "\n",
    "Let's calculate the loss for the current untrained model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tMAT4DcMPwI-"
   },
   "outputs": [],
   "source": [
    "let untrainedLogits = model(firstTrainFeatures)\n",
    "let untrainedLoss = softmaxCrossEntropy(logits: untrainedLogits, labels: firstTrainLabels)\n",
    "print(\"Loss test: \\(untrainedLoss)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lOxFimtlKruu"
   },
   "source": [
    "### Create an optimizer\n",
    "\n",
    "An *[optimizer](https://developers.google.com/machine-learning/crash-course/glossary#optimizer)* applies the computed gradients to the model's variables to minimize the `loss` function. You can think of the loss function as a curved surface (see Figure 3) and we want to find its lowest point by walking around. The gradients point in the direction of steepest ascent—so we'll travel the opposite way and move down the hill. By iteratively calculating the loss and gradient for each batch, we'll adjust the model during training. Gradually, the model will find the best combination of weights and bias to minimize loss. And the lower the loss, the better the model's predictions.\n",
    "\n",
    "<table>\n",
    "  <tr><td>\n",
    "    <img src=\"https://cs231n.github.io/assets/nn3/opt1.gif\" width=\"70%\"\n",
    "         alt=\"Optimization algorithms visualized over time in 3D space.\">\n",
    "  </td></tr>\n",
    "  <tr><td align=\"center\">\n",
    "    <b>Figure 3.</b> Optimization algorithms visualized over time in 3D space.<br/>(Source: <a href=\"http://cs231n.github.io/neural-networks-3/\">Stanford class CS231n</a>, MIT License, Image credit: <a href=\"https://twitter.com/alecrad\">Alec Radford</a>)\n",
    "  </td></tr>\n",
    "</table>\n",
    "\n",
    "Swift for TensorFlow has many [optimization algorithms](https://github.com/tensorflow/swift-apis/tree/main/Sources/TensorFlow/Optimizers) available for training. This model uses the SGD optimizer that implements the *[stochastic gradient descent](https://developers.google.com/machine-learning/crash-course/glossary#gradient_descent)* (SGD) algorithm. The `learningRate` sets the step size to take for each iteration down the hill. This is a *hyperparameter* that you'll commonly adjust to achieve better results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8xxi2NNGKwG_"
   },
   "outputs": [],
   "source": [
    "let optimizer = SGD(for: model, learningRate: 0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pJVRZ0hP52ZB"
   },
   "source": [
    "Let's use `optimizer` to take a single gradient descent step. First, we compute the gradient of the loss with respect to the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rxRNTFVe56RG"
   },
   "outputs": [],
   "source": [
    "let (loss, grads) = valueWithGradient(at: model) { model -> Tensor<Float> in\n",
    "    let logits = model(firstTrainFeatures)\n",
    "    return softmaxCrossEntropy(logits: logits, labels: firstTrainLabels)\n",
    "}\n",
    "print(\"Current loss: \\(loss)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5B27cIT0O0nE"
   },
   "source": [
    "Next, we pass the gradient that we just calculated to the optimizer, which updates the model's differentiable variables accordingly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "icyvh-o6O0nF"
   },
   "outputs": [],
   "source": [
    "optimizer.update(&model, along: grads)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nhpgM7UpO0nG"
   },
   "source": [
    "If we calculate the loss again, it should be smaller, because gradient descent steps (usually) decrease the loss:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aw0OzyojAa39"
   },
   "outputs": [],
   "source": [
    "let logitsAfterOneStep = model(firstTrainFeatures)\n",
    "let lossAfterOneStep = softmaxCrossEntropy(logits: logitsAfterOneStep, labels: firstTrainLabels)\n",
    "print(\"Next loss: \\(lossAfterOneStep)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7Y2VSELvwAvW"
   },
   "source": [
    "### Training loop\n",
    "\n",
    "With all the pieces in place, the model is ready for training! A training loop feeds the dataset examples into the model to help it make better predictions. The following code block sets up these training steps:\n",
    "\n",
    "1. Iterate over each *epoch*. An epoch is one pass through the dataset.\n",
    "2. Within an epoch, iterate over each batch in the training epoch \n",
    "3. Collate the batch and grab its *features* (`x`) and *label* (`y`).\n",
    "3. Using the collated batch's features, make a prediction and compare it with the label. Measure the inaccuracy of the prediction and use that to calculate the model's loss and gradients.\n",
    "4. Use gradient descent to update the model's variables.\n",
    "5. Keep track of some stats for visualization.\n",
    "6. Repeat for each epoch.\n",
    "\n",
    "The `epochCount` variable is the number of times to loop over the dataset collection. Counter-intuitively, training a model longer does not guarantee a better model. `epochCount` is a *[hyperparameter](https://developers.google.com/machine-learning/glossary/#hyperparameter)* that you can tune. Choosing the right number usually requires both experience and experimentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AIgulGRUhpto"
   },
   "outputs": [],
   "source": [
    "let epochCount = 500\n",
    "var trainAccuracyResults: [Float] = []\n",
    "var trainLossResults: [Float] = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "066kVZQFO0nL"
   },
   "outputs": [],
   "source": [
    "func accuracy(predictions: Tensor<Int32>, truths: Tensor<Int32>) -> Float {\n",
    "    return Tensor<Float>(predictions .== truths).mean().scalarized()\n",
    "}\n",
    "\n",
    "for (epochIndex, epoch) in trainingEpochs.prefix(epochCount).enumerated() {\n",
    "    var epochLoss: Float = 0\n",
    "    var epochAccuracy: Float = 0\n",
    "    var batchCount: Int = 0\n",
    "    for batchSamples in epoch {\n",
    "        let batch = batchSamples.collated\n",
    "        let (loss, grad) = valueWithGradient(at: model) { (model: IrisModel) -> Tensor<Float> in\n",
    "            let logits = model(batch.features)\n",
    "            return softmaxCrossEntropy(logits: logits, labels: batch.labels)\n",
    "        }\n",
    "        optimizer.update(&model, along: grad)\n",
    "        \n",
    "        let logits = model(batch.features)\n",
    "        epochAccuracy += accuracy(predictions: logits.argmax(squeezingAxis: 1), truths: batch.labels)\n",
    "        epochLoss += loss.scalarized()\n",
    "        batchCount += 1\n",
    "    }\n",
    "    epochAccuracy /= Float(batchCount)\n",
    "    epochLoss /= Float(batchCount)\n",
    "    trainAccuracyResults.append(epochAccuracy)\n",
    "    trainLossResults.append(epochLoss)\n",
    "    if epochIndex % 50 == 0 {\n",
    "        print(\"Epoch \\(epochIndex): Loss: \\(epochLoss), Accuracy: \\(epochAccuracy)\")\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2FQHVUnm_rjw"
   },
   "source": [
    "### Visualize the loss function over time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j3wdbmtLVTyr"
   },
   "source": [
    "While it's helpful to print out the model's training progress, it's often *more* helpful to see this progress. We can create basic charts using Python's `matplotlib` module.\n",
    "\n",
    "Interpreting these charts takes some experience, but you really want to see the *loss* go down and the *accuracy* go up."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "agjvNd2iUGFn"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize: [12, 8])\n",
    "\n",
    "let accuracyAxes = plt.subplot(2, 1, 1)\n",
    "accuracyAxes.set_ylabel(\"Accuracy\")\n",
    "accuracyAxes.plot(trainAccuracyResults)\n",
    "\n",
    "let lossAxes = plt.subplot(2, 1, 2)\n",
    "lossAxes.set_ylabel(\"Loss\")\n",
    "lossAxes.set_xlabel(\"Epoch\")\n",
    "lossAxes.plot(trainLossResults)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "axA6WuGVO0nR"
   },
   "source": [
    "Note that the y-axes of the graphs are not zero-based."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Zg8GoMZhLpGH"
   },
   "source": [
    "## Evaluate the model's effectiveness\n",
    "\n",
    "Now that the model is trained, we can get some statistics on its performance.\n",
    "\n",
    "*Evaluating* means determining how effectively the model makes predictions. To determine the model's effectiveness at iris classification, pass some sepal and petal measurements to the model and ask the model to predict what iris species they represent. Then compare the model's prediction against the actual label.  For example, a model that picked the correct species on half the input examples has an *[accuracy](https://developers.google.com/machine-learning/glossary/#accuracy)* of `0.5`. Figure 4 shows a slightly more effective model, getting 4 out of 5 predictions correct at 80% accuracy:\n",
    "\n",
    "<table cellpadding=\"8\" border=\"0\">\n",
    "  <colgroup>\n",
    "    <col span=\"4\" >\n",
    "    <col span=\"1\" bgcolor=\"lightblue\">\n",
    "    <col span=\"1\" bgcolor=\"lightgreen\">\n",
    "  </colgroup>\n",
    "  <tr bgcolor=\"lightgray\">\n",
    "    <th colspan=\"4\">Example features</th>\n",
    "    <th colspan=\"1\">Label</th>\n",
    "    <th colspan=\"1\" >Model prediction</th>\n",
    "  </tr>\n",
    "  <tr>\n",
    "    <td>5.9</td><td>3.0</td><td>4.3</td><td>1.5</td><td align=\"center\">1</td><td align=\"center\">1</td>\n",
    "  </tr>\n",
    "  <tr>\n",
    "    <td>6.9</td><td>3.1</td><td>5.4</td><td>2.1</td><td align=\"center\">2</td><td align=\"center\">2</td>\n",
    "  </tr>\n",
    "  <tr>\n",
    "    <td>5.1</td><td>3.3</td><td>1.7</td><td>0.5</td><td align=\"center\">0</td><td align=\"center\">0</td>\n",
    "  </tr>\n",
    "  <tr>\n",
    "    <td>6.0</td> <td>3.4</td> <td>4.5</td> <td>1.6</td> <td align=\"center\">1</td><td align=\"center\" bgcolor=\"red\">2</td>\n",
    "  </tr>\n",
    "  <tr>\n",
    "    <td>5.5</td><td>2.5</td><td>4.0</td><td>1.3</td><td align=\"center\">1</td><td align=\"center\">1</td>\n",
    "  </tr>\n",
    "  <tr><td align=\"center\" colspan=\"6\">\n",
    "    <b>Figure 4.</b> An iris classifier that is 80% accurate.<br/>&nbsp;\n",
    "  </td></tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z-EvK7hGL0d8"
   },
   "source": [
    "### Setup the test dataset\n",
    "\n",
    "Evaluating the model is similar to training the model. The biggest difference is the examples come from a separate *[test set](https://developers.google.com/machine-learning/crash-course/glossary#test_set)* rather than the training set. To fairly assess a model's effectiveness, the examples used to evaluate a model must be different from the examples used to train the model.\n",
    "\n",
    "The setup for the test dataset is similar to the setup for training dataset. Download the test set from http://download.tensorflow.org/data/iris_test.csv:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SRMWCu30bnxH"
   },
   "outputs": [],
   "source": [
    "let testDataFilename = \"iris_test.csv\"\n",
    "download(from: \"http://download.tensorflow.org/data/iris_test.csv\", to: testDataFilename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jEPPL6FUO0nV"
   },
   "source": [
    " Now load it into a an array of `IrisBatch`es:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "w6SCt95HO0nW"
   },
   "outputs": [],
   "source": [
    "let testDataset = loadIrisDatasetFromCSV(\n",
    "    contentsOf: testDataFilename, hasHeader: true,\n",
    "    featureColumns: [0, 1, 2, 3], labelColumns: [4]).inBatches(of: batchSize)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HFuOKXJdMAdm"
   },
   "source": [
    "### Evaluate the model on the test dataset\n",
    "\n",
    "Unlike the training stage, the model only evaluates a single [epoch](https://developers.google.com/machine-learning/glossary/#epoch) of the test data. In the following code cell, we iterate over each example in the test set and compare the model's prediction against the actual label. This is used to measure the model's accuracy across the entire test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Tj4Rs8gwO0nY"
   },
   "outputs": [],
   "source": [
    "// NOTE: Only a single batch will run in the loop since the batchSize we're using is larger than the test set size\n",
    "for batchSamples in testDataset {\n",
    "    let batch = batchSamples.collated\n",
    "    let logits = model(batch.features)\n",
    "    let predictions = logits.argmax(squeezingAxis: 1)\n",
    "    print(\"Test batch accuracy: \\(accuracy(predictions: predictions, truths: batch.labels))\")\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HcKEZMtCOeK-"
   },
   "source": [
    "We can see on the first batch, for example, the model is usually correct:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uNwt2eMeOane"
   },
   "outputs": [],
   "source": [
    "let firstTestBatch = testDataset.first!.collated\n",
    "let firstTestBatchLogits = model(firstTestBatch.features)\n",
    "let firstTestBatchPredictions = firstTestBatchLogits.argmax(squeezingAxis: 1)\n",
    "\n",
    "print(firstTestBatchPredictions)\n",
    "print(firstTestBatch.labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7Li2r1tYvW7S"
   },
   "source": [
    "## Use the trained model to make predictions\n",
    "\n",
    "We've trained a model and demonstrated that it's good—but not perfect—at classifying iris species. Now let's use the trained model to make some predictions on [unlabeled examples](https://developers.google.com/machine-learning/glossary/#unlabeled_example); that is, on examples that contain features but not a label.\n",
    "\n",
    "In real-life, the unlabeled examples could come from lots of different sources including apps, CSV files, and data feeds. For now, we're going to manually provide three unlabeled examples to predict their labels. Recall, the label numbers are mapped to a named representation as:\n",
    "\n",
    "* `0`: Iris setosa\n",
    "* `1`: Iris versicolor\n",
    "* `2`: Iris virginica"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MTYOZr27O0ne"
   },
   "outputs": [],
   "source": [
    "let unlabeledDataset: Tensor<Float> =\n",
    "    [[5.1, 3.3, 1.7, 0.5],\n",
    "     [5.9, 3.0, 4.2, 1.5],\n",
    "     [6.9, 3.1, 5.4, 2.1]]\n",
    "\n",
    "let unlabeledDatasetPredictions = model(unlabeledDataset)\n",
    "\n",
    "for i in 0..<unlabeledDatasetPredictions.shape[0] {\n",
    "    let logits = unlabeledDatasetPredictions[i]\n",
    "    let classIdx = logits.argmax().scalar!\n",
    "    print(\"Example \\(i) prediction: \\(classNames[Int(classIdx)]) (\\(softmax(logits)))\")\n",
    "}"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "model_training_walkthrough.ipynb",
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Swift",
   "name": "swift"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}