[XLA] Fixup the experimental_compile notebook syntax
PiperOrigin-RevId: 288334343 Change-Id: I1f594f20230be9e20f4a2275362d5c8709c113c2
This commit is contained in:
George Karpenkov
TensorFlower Gardener
parent
2d10fe3cd8
commit
93fa0c9c7a
@ -1,268 +1,283 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "Using XLA with tf.function",
|
||||
"provenance": [],
|
||||
"collapsed_sections": [],
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
}
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "f4TSNCvpENrW"
|
||||
},
|
||||
"source": [
|
||||
"##### Copyright 2019 The TensorFlow Authors."
|
||||
]
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "f4TSNCvpENrW"
|
||||
},
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"##### Copyright 2019 The TensorFlow Authors."
|
||||
]
|
||||
},
|
||||
{
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"colab_type": "code",
|
||||
"id": "vamNSA0vEP-m",
|
||||
"colab": {}
|
||||
},
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"#@title 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."
|
||||
],
|
||||
"execution_count": 0,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "e1oSi4lHFt3z"
|
||||
},
|
||||
"source": [
|
||||
"# Using XLA via `tf.function` and `experimental_compile`"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "sDy5lSBd4BDE",
|
||||
"colab_type": "text"
|
||||
},
|
||||
"source": [
|
||||
"In this colab, we train a TensorFlow model to classify the MNIST dataset, where the training function is compiled using XLA.\n",
|
||||
"\n",
|
||||
"We start by loading TensorFlow, with eager execution enabled."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "b7noD9NjFRL-"
|
||||
},
|
||||
"source": [
|
||||
"<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
|
||||
" <td>\n",
|
||||
" <a target=\"_blank\" href=\"https://www.tensorflow.org/xla/tutorials/xla_compile\"><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/xla_compile.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/xla_compile.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
|
||||
" </td>\n",
|
||||
"</table>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab_type": "code",
|
||||
"id": "45kUPj5ZFrRa"
|
||||
},
|
||||
"source": [
|
||||
"import tensorflow as tf\n",
|
||||
"\n",
|
||||
"tf.enable_eager_execution()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "GZVNiRmTDV-5"
|
||||
},
|
||||
"source": [
|
||||
"Then, we define some necessary constants and prepare the MNIST dataset."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab_type": "code",
|
||||
"id": "f37TSEGvGX4_",
|
||||
"colab": {}
|
||||
},
|
||||
"source": [
|
||||
"# Size of each input image, 28 x 28 pixels\n",
|
||||
"IMAGE_SIZE = 28 * 28\n",
|
||||
"# Number of distinct number labels, [0..9]\n",
|
||||
"NUM_CLASSES = 10\n",
|
||||
"# Number of examples in each training batch (step)\n",
|
||||
"TRAIN_BATCH_SIZE = 100\n",
|
||||
"# Number of training steps to run\n",
|
||||
"TRAIN_STEPS = 1000\n",
|
||||
"\n",
|
||||
"# Loads MNIST dataset.\n",
|
||||
"train, test = tf.keras.datasets.mnist.load_data()\n",
|
||||
"train_ds = tf.data.Dataset.from_tensor_slices(train).batch(TRAIN_BATCH_SIZE).repeat()\n",
|
||||
"\n",
|
||||
"# Casting from raw data to the required datatypes.\n",
|
||||
"def cast(images, labels):\n",
|
||||
" images = tf.cast(\n",
|
||||
" tf.reshape(images, [-1, IMAGE_SIZE]), tf.float32)\n",
|
||||
" labels = tf.cast(labels, tf.int64)\n",
|
||||
" return (images, labels)"
|
||||
],
|
||||
"execution_count": 0,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "lv7I-u_82v1S",
|
||||
"colab_type": "text"
|
||||
},
|
||||
"source": [
|
||||
"Finally, we define the model and the optimizer. For the model, we shall use a single dense layer."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "7O2NcEfG206Q",
|
||||
"colab_type": "code",
|
||||
"colab": {}
|
||||
},
|
||||
"source": [
|
||||
"layer = tf.keras.layers.Dense(NUM_CLASSES)\n",
|
||||
"optimizer = tf.keras.optimizers.Adam()\n"
|
||||
],
|
||||
"execution_count": 0,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "x_ZehpZP-SfS"
|
||||
},
|
||||
"source": [
|
||||
"# Define the training function\n",
|
||||
"\n",
|
||||
"In the training function, we get predicted labels using the layer defined above, and then we minimize the gradient of the loss using the optimizer. In order to compile the computation using XLA, we place it inside `tf.function` with `experimental_compile=True`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab_type": "code",
|
||||
"id": "ZbhJl_WvGa3g",
|
||||
"colab": {}
|
||||
},
|
||||
"source": [
|
||||
"@tf.function(experimental_compile=True)\n",
|
||||
"def train_mnist(images, labels):\n",
|
||||
" images, labels = cast(images, labels)\n",
|
||||
"\n",
|
||||
" with tf.GradientTape() as tape:\n",
|
||||
" predicted_labels = layer(images)\n",
|
||||
" loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
|
||||
" logits=predicted_labels, labels=labels\n",
|
||||
" ))\n",
|
||||
" layer_variables = layer.trainable_variables\n",
|
||||
" grads = tape.gradient(loss, layer_variables)\n",
|
||||
" optimizer.apply_gradients(zip(grads, layer_variables))\n"
|
||||
],
|
||||
"execution_count": 0,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "EZD1m_n1DxAF"
|
||||
},
|
||||
"source": [
|
||||
"# Train and test the model"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "gukC2Hol3sFZ",
|
||||
"colab_type": "text"
|
||||
},
|
||||
"source": [
|
||||
"Once we have defined the training function, we can define the model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab_type": "code",
|
||||
"id": "qe28bAHNHUG2",
|
||||
"colab": {}
|
||||
},
|
||||
"source": [
|
||||
"for images, labels in train_ds:\n",
|
||||
" if optimizer.iterations > TRAIN_STEPS:\n",
|
||||
" break\n",
|
||||
" train_mnist(images, labels)"
|
||||
],
|
||||
"execution_count": 0,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "qgsKmz3n2UiW"
|
||||
},
|
||||
"source": [
|
||||
"And, finally, check the accuracy:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab_type": "code",
|
||||
"id": "_GxF6jTRHVuA"
|
||||
},
|
||||
"source": [
|
||||
"images, labels = cast(test[0], test[1])\n",
|
||||
"predicted_labels = layer(images)\n",
|
||||
"correct_prediction = tf.equal(tf.argmax(predicted_labels, 1), labels)\n",
|
||||
"accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))\n",
|
||||
"print(\"Prediction accuracy after training: %s\" % accuracy)"
|
||||
],
|
||||
"execution_count": 0
|
||||
}
|
||||
]
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"cellView": "form",
|
||||
"colab": {},
|
||||
"colab_type": "code",
|
||||
"id": "vamNSA0vEP-m"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"#@title 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": {
|
||||
"colab_type": "text",
|
||||
"id": "e1oSi4lHFt3z"
|
||||
},
|
||||
"source": [
|
||||
"# Using XLA via `tf.function` and `experimental_compile`"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "sDy5lSBd4BDE"
|
||||
},
|
||||
"source": [
|
||||
"In this colab, we train a TensorFlow model to classify the MNIST dataset, where the training function is compiled using XLA.\n",
|
||||
"\n",
|
||||
"We start by loading TensorFlow, with eager execution enabled."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "b7noD9NjFRL-"
|
||||
},
|
||||
"source": [
|
||||
"<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
|
||||
" <td>\n",
|
||||
" <a target=\"_blank\" href=\"https://www.tensorflow.org/xla/tutorials/experimental_compile\"><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=\"../g3doc/tutorials/experimental_compile.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/experimental_compile.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
|
||||
" </td>\n",
|
||||
"</table>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"colab_type": "code",
|
||||
"id": "45kUPj5ZFrRa"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import tensorflow as tf\n",
|
||||
"\n",
|
||||
"tf.enable_eager_execution()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "GZVNiRmTDV-5"
|
||||
},
|
||||
"source": [
|
||||
"Then, we define some necessary constants and prepare the MNIST dataset."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"metadata": {
|
||||
"colab": {},
|
||||
"colab_type": "code",
|
||||
"id": "f37TSEGvGX4_"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Size of each input image, 28 x 28 pixels\n",
|
||||
"IMAGE_SIZE = 28 * 28\n",
|
||||
"# Number of distinct number labels, [0..9]\n",
|
||||
"NUM_CLASSES = 10\n",
|
||||
"# Number of examples in each training batch (step)\n",
|
||||
"TRAIN_BATCH_SIZE = 100\n",
|
||||
"# Number of training steps to run\n",
|
||||
"TRAIN_STEPS = 1000\n",
|
||||
"\n",
|
||||
"# Loads MNIST dataset.\n",
|
||||
"train, test = tf.keras.datasets.mnist.load_data()\n",
|
||||
"train_ds = tf.data.Dataset.from_tensor_slices(train).batch(TRAIN_BATCH_SIZE).repeat()\n",
|
||||
"\n",
|
||||
"# Casting from raw data to the required datatypes.\n",
|
||||
"def cast(images, labels):\n",
|
||||
" images = tf.cast(\n",
|
||||
" tf.reshape(images, [-1, IMAGE_SIZE]), tf.float32)\n",
|
||||
" labels = tf.cast(labels, tf.int64)\n",
|
||||
" return (images, labels)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "lv7I-u_82v1S"
|
||||
},
|
||||
"source": [
|
||||
"Finally, we define the model and the optimizer. For the model, we shall use a single dense layer."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 4,
|
||||
"metadata": {
|
||||
"colab": {},
|
||||
"colab_type": "code",
|
||||
"id": "7O2NcEfG206Q"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"layer = tf.keras.layers.Dense(NUM_CLASSES)\n",
|
||||
"optimizer = tf.keras.optimizers.Adam()\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "x_ZehpZP-SfS"
|
||||
},
|
||||
"source": [
|
||||
"# Define the training function\n",
|
||||
"\n",
|
||||
"In the training function, we get predicted labels using the layer defined above, and then we minimize the gradient of the loss using the optimizer. In order to compile the computation using XLA, we place it inside `tf.function` with `experimental_compile=True`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 5,
|
||||
"metadata": {
|
||||
"colab": {},
|
||||
"colab_type": "code",
|
||||
"id": "ZbhJl_WvGa3g"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"@tf.function(experimental_compile=True)\n",
|
||||
"def train_mnist(images, labels):\n",
|
||||
" images, labels = cast(images, labels)\n",
|
||||
"\n",
|
||||
" with tf.GradientTape() as tape:\n",
|
||||
" predicted_labels = layer(images)\n",
|
||||
" loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
|
||||
" logits=predicted_labels, labels=labels\n",
|
||||
" ))\n",
|
||||
" layer_variables = layer.trainable_variables\n",
|
||||
" grads = tape.gradient(loss, layer_variables)\n",
|
||||
" optimizer.apply_gradients(zip(grads, layer_variables))\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "EZD1m_n1DxAF"
|
||||
},
|
||||
"source": [
|
||||
"# Train and test the model"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "gukC2Hol3sFZ"
|
||||
},
|
||||
"source": [
|
||||
"Once we have defined the training function, we can define the model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"metadata": {
|
||||
"colab": {},
|
||||
"colab_type": "code",
|
||||
"id": "qe28bAHNHUG2"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"for images, labels in train_ds:\n",
|
||||
" if optimizer.iterations > TRAIN_STEPS:\n",
|
||||
" break\n",
|
||||
" train_mnist(images, labels)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"colab_type": "text",
|
||||
"id": "qgsKmz3n2UiW"
|
||||
},
|
||||
"source": [
|
||||
"And, finally, check the accuracy:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {
|
||||
"colab_type": "code",
|
||||
"id": "_GxF6jTRHVuA"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"images, labels = cast(test[0], test[1])\n",
|
||||
"predicted_labels = layer(images)\n",
|
||||
"correct_prediction = tf.equal(tf.argmax(predicted_labels, 1), labels)\n",
|
||||
"accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))\n",
|
||||
"print(\"Prediction accuracy after training: %s\" % accuracy)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"collapsed_sections": [],
|
||||
"name": "Using XLA with tf.function",
|
||||
"provenance": [],
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.7.5rc1"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user