<commit_before>{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "CIFT with XLA.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "metadata": {
        "colab_type": "text",
        "id": "b7noD9NjFRL-"
      },
      "cell_type": "markdown",
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/xla/tutorials/autoclustering_xla\"><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/tensorflow/blob/master/tensorflow/compiler/xla/g3doc/tutorials/autoclustering_xla.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/tensorflow/blob/master/tensorflow/compiler/xla/g3doc/tutorials/autoclustering_xla.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": {
        "colab_type": "text",
        "id": "mz65veHXsmnS"
      },
      "source": [
        "# Classifying CIFAR-10 with XLA\n",
        "\n",
        "This tutorial trains a TensorFlow model to classify the [CIFAR-10](https://en.wikipedia.org/wiki/CIFAR-10) dataset, and we compile it using XLA.\n",
        "\n",
        "Load and normalize the dataset using the Keras API:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "7vm2QsMisCxI"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "# Check that GPU is available: cf. https://colab.research.google.com/notebooks/gpu.ipynb\n",
        "assert(tf.test.is_gpu_available())\n",
        "\n",
        "tf.keras.backend.clear_session()\n",
        "tf.config.optimizer.set_jit(False) # Start with XLA disabled.\n",
        "\n",
        "def load_data():\n",
        "  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()\n",
        "  x_train = x_train.astype('float32') / 256\n",
        "  x_test = x_test.astype('float32') / 256\n",
        "\n",
        "  # Convert class vectors to binary class matrices.\n",
        "  y_train = tf.keras.utils.to_categorical(y_train, num_classes=10)\n",
        "  y_test = tf.keras.utils.to_categorical(y_test, num_classes=10)\n",
        "  return ((x_train, y_train), (x_test, y_test))\n",
        "\n",
        "(x_train, y_train), (x_test, y_test) = load_data()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MgNM2tbgtScx"
      },
      "source": [
        "We define the model, adapted from the Keras [CIFAR-10 example](https://keras.io/examples/cifar10_cnn/):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3ZRQSwoRsKM_"
      },
      "outputs": [],
      "source": [
        "def generate_model():\n",
        "  return tf.keras.models.Sequential([\n",
        "    tf.keras.layers.Conv2D(32, (3, 3), padding='same', input_shape=x_train.shape[1:]),\n",
        "    tf.keras.layers.Activation('relu'),\n",
        "    tf.keras.layers.Conv2D(32, (3, 3)),\n",
        "    tf.keras.layers.Activation('relu'),\n",
        "    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
        "    tf.keras.layers.Dropout(0.25),\n",
        "\n",
        "    tf.keras.layers.Conv2D(64, (3, 3), padding='same'),\n",
        "    tf.keras.layers.Activation('relu'),\n",
        "    tf.keras.layers.Conv2D(64, (3, 3)),\n",
        "    tf.keras.layers.Activation('relu'),\n",
        "    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
        "    tf.keras.layers.Dropout(0.25),\n",
        "\n",
        "    tf.keras.layers.Flatten(),\n",
        "    tf.keras.layers.Dense(512),\n",
        "    tf.keras.layers.Activation('relu'),\n",
        "    tf.keras.layers.Dropout(0.5),\n",
        "    tf.keras.layers.Dense(10),\n",
        "    tf.keras.layers.Activation('softmax')\n",
        "  ])\n",
        "\n",
        "model = generate_model()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-M4GtGDZtb8a"
      },
      "source": [
        "We train the model using the\n",
        "[RMSprop](https://www.tensorflow.org/api_docs/python/tf/train/RMSPropOptimizer)\n",
        "optimizer:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "UKCmrhF0tiMa"
      },
      "outputs": [],
      "source": [
        "def compile_model(model):\n",
        "  opt = tf.keras.optimizers.RMSprop(lr=0.0001, decay=1e-6)\n",
        "  model.compile(loss='categorical_crossentropy',\n",
        "                optimizer=opt,\n",
        "                metrics=['accuracy'])\n",
        "  return model\n",
        "\n",
        "model = compile_model(model)\n",
        "\n",
        "def train_model(model, x_train, y_train, x_test, y_test, epochs=25):\n",
        "  model.fit(x_train, y_train, batch_size=256, epochs=epochs, validation_data=(x_test, y_test), shuffle=True)\n",
        "\n",
        "def warmup(model, x_train, y_train, x_test, y_test):\n",
        "  # Warm up the JIT, we do not wish to measure the compilation time.\n",
        "  initial_weights = model.get_weights()\n",
        "  train_model(model, x_train, y_train, x_test, y_test, epochs=1)\n",
        "  model.set_weights(initial_weights)\n",
        "\n",
        "warmup(model, x_train, y_train, x_test, y_test)\n",
        "train_model(model, x_train, y_train, x_test, y_test)\n",
        "\n",
        "scores = model.evaluate(x_test, y_test, verbose=1)\n",
        "print('Test loss:', scores[0])\n",
        "print('Test accuracy:', scores[1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SLpfQ0StRgsu"
      },
      "source": [
        "Now let's train the model again, using the XLA compiler.\n",
        "To enable the compiler in the middle of the application, we need to reset the Keras session."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "jxU-Tzy4SX7p"
      },
      "outputs": [],
      "source": [
        "# We need to clear the session to enable JIT in the middle of the program.\n",
        "tf.keras.backend.clear_session()\n",
        "tf.config.optimizer.set_jit(True) # Enable XLA.\n",
        "model = compile_model(generate_model())\n",<commit_msg>Updating the XLA autoclustering tutorial with the new way to check for GPUs and also added a printout of time in the non-XLA case to compare against the XLA case.<commit_after>{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "CIFT with XLA.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "metadata": {
        "colab_type": "text",
        "id": "b7noD9NjFRL-"
      },
      "cell_type": "markdown",
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/xla/tutorials/autoclustering_xla\"><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/tensorflow/blob/master/tensorflow/compiler/xla/g3doc/tutorials/autoclustering_xla.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/tensorflow/blob/master/tensorflow/compiler/xla/g3doc/tutorials/autoclustering_xla.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": {
        "colab_type": "text",
        "id": "mz65veHXsmnS"
      },
      "source": [
        "# Classifying CIFAR-10 with XLA\n",
        "\n",
        "This tutorial trains a TensorFlow model to classify the [CIFAR-10](https://en.wikipedia.org/wiki/CIFAR-10) dataset, and we compile it using XLA.\n",
        "\n",
        "Load and normalize the dataset using the Keras API:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "7vm2QsMisCxI"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "# Check that GPU is available: cf. https://colab.research.google.com/notebooks/gpu.ipynb\n",
        "assert(tf.test.gpu_device_name())\n",
        "\n",
        "tf.keras.backend.clear_session()\n",
        "tf.config.optimizer.set_jit(False) # Start with XLA disabled.\n",
        "\n",
        "def load_data():\n",
        "  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()\n",
        "  x_train = x_train.astype('float32') / 256\n",
        "  x_test = x_test.astype('float32') / 256\n",
        "\n",
        "  # Convert class vectors to binary class matrices.\n",
        "  y_train = tf.keras.utils.to_categorical(y_train, num_classes=10)\n",
        "  y_test = tf.keras.utils.to_categorical(y_test, num_classes=10)\n",
        "  return ((x_train, y_train), (x_test, y_test))\n",
        "\n",
        "(x_train, y_train), (x_test, y_test) = load_data()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MgNM2tbgtScx"
      },
      "source": [
        "We define the model, adapted from the Keras [CIFAR-10 example](https://keras.io/examples/cifar10_cnn/):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3ZRQSwoRsKM_"
      },
      "outputs": [],
      "source": [
        "def generate_model():\n",
        "  return tf.keras.models.Sequential([\n",
        "    tf.keras.layers.Conv2D(32, (3, 3), padding='same', input_shape=x_train.shape[1:]),\n",
        "    tf.keras.layers.Activation('relu'),\n",
        "    tf.keras.layers.Conv2D(32, (3, 3)),\n",
        "    tf.keras.layers.Activation('relu'),\n",
        "    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
        "    tf.keras.layers.Dropout(0.25),\n",
        "\n",
        "    tf.keras.layers.Conv2D(64, (3, 3), padding='same'),\n",
        "    tf.keras.layers.Activation('relu'),\n",
        "    tf.keras.layers.Conv2D(64, (3, 3)),\n",
        "    tf.keras.layers.Activation('relu'),\n",
        "    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
        "    tf.keras.layers.Dropout(0.25),\n",
        "\n",
        "    tf.keras.layers.Flatten(),\n",
        "    tf.keras.layers.Dense(512),\n",
        "    tf.keras.layers.Activation('relu'),\n",
        "    tf.keras.layers.Dropout(0.5),\n",
        "    tf.keras.layers.Dense(10),\n",
        "    tf.keras.layers.Activation('softmax')\n",
        "  ])\n",
        "\n",
        "model = generate_model()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-M4GtGDZtb8a"
      },
      "source": [
        "We train the model using the\n",
        "[RMSprop](https://www.tensorflow.org/api_docs/python/tf/train/RMSPropOptimizer)\n",
        "optimizer:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "UKCmrhF0tiMa"
      },
      "outputs": [],
      "source": [
        "def compile_model(model):\n",
        "  opt = tf.keras.optimizers.RMSprop(lr=0.0001, decay=1e-6)\n",
        "  model.compile(loss='categorical_crossentropy',\n",
        "                optimizer=opt,\n",
        "                metrics=['accuracy'])\n",
        "  return model\n",
        "\n",
        "model = compile_model(model)\n",
        "\n",
        "def train_model(model, x_train, y_train, x_test, y_test, epochs=25):\n",
        "  model.fit(x_train, y_train, batch_size=256, epochs=epochs, validation_data=(x_test, y_test), shuffle=True)\n",
        "\n",
        "def warmup(model, x_train, y_train, x_test, y_test):\n",
        "  # Warm up the JIT, we do not wish to measure the compilation time.\n",
        "  initial_weights = model.get_weights()\n",
        "  train_model(model, x_train, y_train, x_test, y_test, epochs=1)\n",
        "  model.set_weights(initial_weights)\n",
        "\n",
        "warmup(model, x_train, y_train, x_test, y_test)\n",
        "%time train_model(model, x_train, y_train, x_test, y_test)\n",
        "\n",
        "scores = model.evaluate(x_test, y_test, verbose=1)\n",
        "print('Test loss:', scores[0])\n",
        "print('Test accuracy:', scores[1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SLpfQ0StRgsu"
      },
      "source": [
        "Now let's train the model again, using the XLA compiler.\n",
        "To enable the compiler in the middle of the application, we need to reset the Keras session."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "jxU-Tzy4SX7p"
      },
      "outputs": [],
      "source": [
        "# We need to clear the session to enable JIT in the middle of the program.\n",
        "tf.keras.backend.clear_session()\n",
        "tf.config.optimizer.set_jit(True) # Enable XLA.\n",
        "model = compile_model(generate_model())\n",
{'fraction_non_alphanumeric': 0.19015222955678296, 'fraction_numerical': 0.015492388522160852, 'mean_word_length': 2.0911930043722675, 'pattern_counts': {'":': 142, '<': 37, '<?xml version=': 0, '>': 37, 'https://': 20, 'lorem ipsum': 0, 'www.': 10, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7320710', 'n_tokens_mistral': 5621, 'n_tokens_neox': 5423, 'n_words': 877}