{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "Tce3stUlHN0L" }, "source": [ "##### Copyright 2024 The TensorFlow Authors.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "tuOe1ymfHZPu" }, "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": { "id": "MfBg1C5NB3X0" }, "source": [ "# Use TPUs\n", "\n", "\n", " \n", " \n", " \n", " \n", "
\n", " View on TensorFlow.org\n", " \n", " Run in Google Colab\n", " \n", " View source on GitHub\n", " \n", " Download notebook\n", "
" ] }, { "cell_type": "markdown", "metadata": { "id": "Ys81cOhXOWUP" }, "source": [ "This guide demonstrates how to perform basic training on [Tensor Processing Units (TPUs)](https://cloud.google.com/tpu/) and TPU Pods, a collection of TPU devices connected by dedicated high-speed network interfaces, with `tf.keras` and custom training loops.\n", "\n", "TPUs are Google's custom-developed application-specific integrated circuits (ASICs) used to accelerate machine learning workloads. They are available through [Google Colab](https://colab.research.google.com/), the [TPU Research Cloud](https://sites.research.google/trc/), and [Cloud TPU](https://cloud.google.com/tpu)." ] }, { "cell_type": "markdown", "metadata": { "id": "ek5Hop74NVKm" }, "source": [ "## Setup" ] }, { "cell_type": "markdown", "metadata": { "id": "ebf7f8489bb7" }, "source": [ "Before you run this Colab notebook, make sure that your hardware accelerator is a TPU by checking your notebook settings: **Runtime** > **Change runtime type** > **Hardware accelerator** > **TPU v2**.\n", "\n", "Import some necessary libraries, including TensorFlow Datasets:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Cw0WRaChRxTL" }, "outputs": [], "source": [ "import tensorflow as tf\n", "\n", "import os\n", "import tensorflow_datasets as tfds" ] }, { "cell_type": "markdown", "metadata": { "id": "yDWaRxSpwBN1" }, "source": [ "## TPU initialization\n", "\n", "TPUs are typically [Cloud TPU](https://cloud.google.com/tpu/docs/) workers, which are different from the local process running the user's Python program. Thus, you need to do some initialization work to connect to the remote cluster and initialize the TPUs. Note that the `tpu` argument to `tf.distribute.cluster_resolver.TPUClusterResolver` is a special address just for Colab. If you are running your code on Google Compute Engine (GCE), you should instead pass in the name of your Cloud TPU." ] }, { "cell_type": "markdown", "metadata": { "id": "dCqWMqvtwOLs" }, "source": [ "Note: The TPU initialization code has to be at the beginning of your program." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "dKPqF8d1wJCV" }, "outputs": [], "source": [ "resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='local')\n", "tf.config.experimental_connect_to_cluster(resolver)\n", "# This is the TPU initialization code that has to be at the beginning.\n", "tf.tpu.experimental.initialize_tpu_system(resolver)\n", "print(\"All devices: \", tf.config.list_logical_devices('TPU'))" ] }, { "cell_type": "markdown", "metadata": { "id": "Mv7kehTZ1Lq_" }, "source": [ "## Manual device placement\n", "\n", "After the TPU is initialized, you can use manual device placement to place the computation on a single TPU device:\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "XRZ4kMoxBNND" }, "outputs": [], "source": [ "a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])\n", "b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])\n", "\n", "with tf.device('/TPU:0'):\n", " c = tf.matmul(a, b)\n", "\n", "print(\"c device: \", c.device)\n", "print(c)" ] }, { "cell_type": "markdown", "metadata": { "id": "_NJm-kgFO0cC" }, "source": [ "## Distribution strategies\n", "\n", "Usually, you run your model on multiple TPUs in a data-parallel way. To distribute your model on multiple TPUs (as well as multiple GPUs or multiple machines), TensorFlow offers the `tf.distribute.Strategy` API. You can replace your distribution strategy and the model will run on any given (TPU) device. Learn more in the [Distributed training with TensorFlow](./distributed_training.ipynb) guide." ] }, { "cell_type": "markdown", "metadata": { "id": "DcDPMZs-9uLJ" }, "source": [ "Using the `tf.distribute.TPUStrategy` option implements synchronous distributed training. TPUs provide their own implementation of efficient all-reduce and other collective operations across multiple TPU cores, which are used in `TPUStrategy`.\n", "\n", "To demonstrate this, create a `tf.distribute.TPUStrategy` object:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "7SO23K8oRpjI" }, "outputs": [], "source": [ "strategy = tf.distribute.TPUStrategy(resolver)" ] }, { "cell_type": "markdown", "metadata": { "id": "JlaAmswWPsU6" }, "source": [ "To replicate a computation so it can run in all TPU cores, you can pass it into the `Strategy.run` API. Below is an example that shows all cores receiving the same inputs `(a, b)` and performing matrix multiplication on each core independently. The outputs will be the values from all the replicas." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "-90CL5uFPTOa" }, "outputs": [], "source": [ "@tf.function\n", "def matmul_fn(x, y):\n", " z = tf.matmul(x, y)\n", " return z\n", "\n", "z = strategy.run(matmul_fn, args=(a, b))\n", "print(z)" ] }, { "cell_type": "markdown", "metadata": { "id": "uxgYl6kGHJLc" }, "source": [ "## Classification on TPUs\n", "\n", "Having covered the basic concepts, consider a more concrete example. This section demonstrates how to use the distribution strategy—`tf.distribute.TPUStrategy`—to train a Keras model on a Cloud TPU." ] }, { "cell_type": "markdown", "metadata": { "id": "gKRALGgt_kCo" }, "source": [ "### Define a Keras model\n", "\n", "Start with a definition of a [`Sequential` Keras model](https://www.tensorflow.org/guide/keras/sequential_model) for image classification on the MNIST dataset. It's no different than what you would use if you were training on CPUs or GPUs. Note that Keras model creation needs to be inside the `Strategy.scope`, so the variables can be created on each TPU device. Other parts of the code are not necessary to be inside the `Strategy` scope." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "DiBiN-Z_R7P7" }, "outputs": [], "source": [ "def create_model():\n", " regularizer = tf.keras.regularizers.L2(1e-5)\n", " return tf.keras.Sequential(\n", " [tf.keras.layers.Conv2D(256, 3, input_shape=(28, 28, 1),\n", " activation='relu',\n", " kernel_regularizer=regularizer),\n", " tf.keras.layers.Conv2D(256, 3,\n", " activation='relu',\n", " kernel_regularizer=regularizer),\n", " tf.keras.layers.Flatten(),\n", " tf.keras.layers.Dense(256,\n", " activation='relu',\n", " kernel_regularizer=regularizer),\n", " tf.keras.layers.Dense(128,\n", " activation='relu',\n", " kernel_regularizer=regularizer),\n", " tf.keras.layers.Dense(10,\n", " kernel_regularizer=regularizer)])" ] }, { "cell_type": "markdown", "metadata": { "id": "h-2qaXgfyONQ" }, "source": [ "This model puts L2 regularization terms on the weights of each layer, so that the custom training loop below can show how you pick them up from `Model.losses`." ] }, { "cell_type": "markdown", "metadata": { "id": "qYOYjYTg_31l" }, "source": [ "### Load the dataset\n", "\n", "Efficient use of the `tf.data.Dataset` API is critical when using a Cloud TPU. You can learn more about dataset performance in the [Input pipeline performance guide](./data_performance.ipynb).\n", "\n", "If you are using [TPU Nodes](https://cloud.google.com/tpu/docs/managing-tpus-tpu-vm), you need to store all data files read by the TensorFlow `Dataset` in [Google Cloud Storage (GCS) buckets](https://cloud.google.com/tpu/docs/storage-buckets). If you are using [TPU VMs](https://cloud.google.com/tpu/docs/users-guide-tpu-vm), you can store data wherever you like. For more information on TPU Nodes and TPU VMs, refer to the [TPU System Architecture](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm) documentation.\n", "\n", "For most use cases, it is recommended to convert your data into the `TFRecord` format and use a `tf.data.TFRecordDataset` to read it. Check the [TFRecord and tf.Example tutorial](../tutorials/load_data/tfrecord.ipynb) for details on how to do this. It is not a hard requirement and you can use other dataset readers, such as `tf.data.FixedLengthRecordDataset` or `tf.data.TextLineDataset`.\n", "\n", "You can load entire small datasets into memory using `tf.data.Dataset.cache`.\n", "\n", "Regardless of the data format used, it is strongly recommended that you use large files on the order of 100MB. This is especially important in this networked setting, as the overhead of opening a file is significantly higher.\n", "\n", "As shown in the code below, you should use the Tensorflow Datasets `tfds.load` module to get a copy of the MNIST training and test data. Note that `try_gcs` is specified to use a copy that is available in a public GCS bucket. If you don't specify this, the TPU will not be able to access the downloaded data." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "noAd416KSCo7" }, "outputs": [], "source": [ "def get_dataset(batch_size, is_training=True):\n", " split = 'train' if is_training else 'test'\n", " dataset, info = tfds.load(name='mnist', split=split, with_info=True,\n", " as_supervised=True, try_gcs=True)\n", "\n", " # Normalize the input data.\n", " def scale(image, label):\n", " image = tf.cast(image, tf.float32)\n", " image /= 255.0\n", " return image, label\n", "\n", " dataset = dataset.map(scale)\n", "\n", " # Only shuffle and repeat the dataset in training. The advantage of having an\n", " # infinite dataset for training is to avoid the potential last partial batch\n", " # in each epoch, so that you don't need to think about scaling the gradients\n", " # based on the actual batch size.\n", " if is_training:\n", " dataset = dataset.shuffle(10000)\n", " dataset = dataset.repeat()\n", "\n", " dataset = dataset.batch(batch_size)\n", "\n", " return dataset" ] }, { "cell_type": "markdown", "metadata": { "id": "mgUC6A-zCMEr" }, "source": [ "### Train the model using Keras high-level APIs\n", "\n", "You can train your model with Keras `Model.fit` and `Model.compile` APIs. There is nothing TPU-specific in this step—you write the code as if you were using multiple GPUs and a `MirroredStrategy` instead of the `TPUStrategy`. You can learn more in the [Distributed training with Keras](../tutorials/distribute/keras.ipynb) tutorial." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "ubmDchPqSIx0" }, "outputs": [], "source": [ "with strategy.scope():\n", " model = create_model()\n", " model.compile(optimizer='adam',\n", " loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n", " metrics=['sparse_categorical_accuracy'])\n", "\n", "batch_size = 200\n", "steps_per_epoch = 60000 // batch_size\n", "validation_steps = 10000 // batch_size\n", "\n", "train_dataset = get_dataset(batch_size, is_training=True)\n", "test_dataset = get_dataset(batch_size, is_training=False)\n", "\n", "model.fit(train_dataset,\n", " epochs=5,\n", " steps_per_epoch=steps_per_epoch,\n", " validation_data=test_dataset,\n", " validation_steps=validation_steps)" ] }, { "cell_type": "markdown", "metadata": { "id": "8hSGBIYtUugJ" }, "source": [ "To reduce Python overhead and maximize the performance of your TPU, pass in the `steps_per_execution` argument to Keras `Model.compile`. In this example, it increases throughput by about 50%:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "M6e3aVVLUorL" }, "outputs": [], "source": [ "with strategy.scope():\n", " model = create_model()\n", " model.compile(optimizer='adam',\n", " # Anything between 2 and `steps_per_epoch` could help here.\n", " steps_per_execution = 50,\n", " loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n", " metrics=['sparse_categorical_accuracy'])\n", "\n", "model.fit(train_dataset,\n", " epochs=5,\n", " steps_per_epoch=steps_per_epoch,\n", " validation_data=test_dataset,\n", " validation_steps=validation_steps)" ] }, { "cell_type": "markdown", "metadata": { "id": "0rRALBZNCO4A" }, "source": [ "### Train the model using a custom training loop\n", "\n", "You can also create and train your model using `tf.function` and `tf.distribute` APIs directly. You can use the `Strategy.distribute_datasets_from_function` API to distribute the `tf.data.Dataset` given a dataset function. Note that in the example below the batch size passed into the `Dataset` is the per-replica batch size instead of the global batch size. To learn more, check out the [Custom training with `tf.distribute.Strategy`](../tutorials/distribute/custom_training.ipynb) tutorial.\n" ] }, { "cell_type": "markdown", "metadata": { "id": "DxdgXPAL6iFE" }, "source": [ "First, create the model, datasets and `tf.function`s:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "9aHhqwao2Fxi" }, "outputs": [], "source": [ "# Create the model, optimizer and metrics inside the `tf.distribute.Strategy`\n", "# scope, so that the variables can be mirrored on each device.\n", "with strategy.scope():\n", " model = create_model()\n", " optimizer = tf.keras.optimizers.Adam()\n", " training_loss = tf.keras.metrics.Mean('training_loss', dtype=tf.float32)\n", " training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n", " 'training_accuracy', dtype=tf.float32)\n", "\n", "# Calculate per replica batch size, and distribute the `tf.data.Dataset`s\n", "# on each TPU worker.\n", "per_replica_batch_size = batch_size // strategy.num_replicas_in_sync\n", "\n", "train_dataset = strategy.distribute_datasets_from_function(\n", " lambda _: get_dataset(per_replica_batch_size, is_training=True))\n", "\n", "@tf.function\n", "def train_step(iterator):\n", " \"\"\"The step function for one training step.\"\"\"\n", "\n", " def step_fn(inputs):\n", " \"\"\"The computation to run on each TPU device.\"\"\"\n", " images, labels = inputs\n", " with tf.GradientTape() as tape:\n", " logits = model(images, training=True)\n", " per_example_loss = tf.keras.losses.sparse_categorical_crossentropy(\n", " labels, logits, from_logits=True)\n", " loss = tf.nn.compute_average_loss(per_example_loss)\n", " model_losses = model.losses\n", " if model_losses:\n", " loss += tf.nn.scale_regularization_loss(tf.add_n(model_losses))\n", "\n", " grads = tape.gradient(loss, model.trainable_variables)\n", " optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))\n", " training_loss.update_state(loss * strategy.num_replicas_in_sync)\n", " training_accuracy.update_state(labels, logits)\n", "\n", " strategy.run(step_fn, args=(next(iterator),))" ] }, { "cell_type": "markdown", "metadata": { "id": "Ibi7Z97V6xsQ" }, "source": [ "Then, run the training loop:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "1du5cXWt6Vtw" }, "outputs": [], "source": [ "steps_per_eval = 10000 // batch_size\n", "\n", "train_iterator = iter(train_dataset)\n", "for epoch in range(5):\n", " print('Epoch: {}/5'.format(epoch))\n", "\n", " for step in range(steps_per_epoch):\n", " train_step(train_iterator)\n", " print('Current step: {}, training loss: {}, training accuracy: {}%'.format(\n", " optimizer.iterations.numpy(),\n", " round(float(training_loss.result()), 4),\n", " round(float(training_accuracy.result()) * 100, 2)))\n", " training_loss.reset_states()\n", " training_accuracy.reset_states()" ] }, { "cell_type": "markdown", "metadata": { "id": "TnZJUM3qIjKu" }, "source": [ "### Improving performance with multiple steps inside `tf.function`\n", "\n", "You can improve the performance by running multiple steps within a `tf.function`. This is achieved by wrapping the `Strategy.run` call with a `tf.range` inside `tf.function`, and AutoGraph will convert it to a `tf.while_loop` on the TPU worker. You can learn more about `tf.function`s in the [Better performance with `tf.function`](./function.ipynb) guide.\n", "\n", "Despite the improved performance, there are tradeoffs with this method compared to running a single step inside a `tf.function`. Running multiple steps in a `tf.function` is less flexible—you cannot run things eagerly or arbitrary Python code within the steps.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "2grYvXLzJYkP" }, "outputs": [], "source": [ "@tf.function\n", "def train_multiple_steps(iterator, steps):\n", " \"\"\"The step function for one training step.\"\"\"\n", "\n", " def step_fn(inputs):\n", " \"\"\"The computation to run on each TPU device.\"\"\"\n", " images, labels = inputs\n", " with tf.GradientTape() as tape:\n", " logits = model(images, training=True)\n", " per_example_loss = tf.keras.losses.sparse_categorical_crossentropy(\n", " labels, logits, from_logits=True)\n", " loss = tf.nn.compute_average_loss(per_example_loss)\n", " model_losses = model.losses\n", " if model_losses:\n", " loss += tf.nn.scale_regularization_loss(tf.add_n(model_losses))\n", " grads = tape.gradient(loss, model.trainable_variables)\n", " optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))\n", " training_loss.update_state(loss * strategy.num_replicas_in_sync)\n", " training_accuracy.update_state(labels, logits)\n", "\n", " for _ in tf.range(steps):\n", " strategy.run(step_fn, args=(next(iterator),))\n", "\n", "# Convert `steps_per_epoch` to `tf.Tensor` so the `tf.function` won't get\n", "# retraced if the value changes.\n", "train_multiple_steps(train_iterator, tf.convert_to_tensor(steps_per_epoch))\n", "\n", "print('Current step: {}, training loss: {}, training accuracy: {}%'.format(\n", " optimizer.iterations.numpy(),\n", " round(float(training_loss.result()), 4),\n", " round(float(training_accuracy.result()) * 100, 2)))" ] }, { "cell_type": "markdown", "metadata": { "id": "WBKVhMvWjibf" }, "source": [ "## Next steps\n", "\n", "To learn more about Cloud TPUs and how to use them:\n", "\n", "- [Google Cloud TPU](https://cloud.google.com/tpu): The Google Cloud TPU homepage.\n", "- [Google Cloud TPU documentation](https://cloud.google.com/tpu/docs/): Google Cloud TPU documentation, which includes:\n", " - [Introduction to Cloud TPU](https://cloud.google.com/tpu/docs/intro-to-tpu): An overview of working with Cloud TPUs.\n", " - [Cloud TPU quickstarts](https://cloud.google.com/tpu/docs/quick-starts): Quickstart introductions to working with Cloud TPU VMs using TensorFlow and other main machine learning frameworks.\n", "- [Google Cloud TPU Colab notebooks](https://cloud.google.com/tpu/docs/colabs): End-to-end training examples.\n", "- [Google Cloud TPU performance guide](https://cloud.google.com/tpu/docs/performance-guide): Enhance Cloud TPU performance further by adjusting Cloud TPU configuration parameters for your application\n", "- [Distributed training with TensorFlow](./distributed_training.ipynb): How to use distribution strategies—including `tf.distribute.TPUStrategy`—with examples showing best practices.\n", "- TPU embeddings: TensorFlow includes specialized support for training embeddings on TPUs via `tf.tpu.experimental.embedding`. In addition, [TensorFlow Recommenders](https://www.tensorflow.org/recommenders) has `tfrs.layers.embedding.TPUEmbedding`. Embeddings provide efficient and dense representations, capturing complex similarities and relationships between features. TensorFlow's TPU-specific embedding support allows you to train embeddings that are larger than the memory of a single TPU device, and to use sparse and ragged inputs on TPUs.\n", "- [TPU Research Cloud (TRC)](https://sites.research.google/trc/about/): TRC enables researchers to apply for access to a cluster of more than 1,000 Cloud TPU devices.\n" ] } ], "metadata": { "accelerator": "TPU", "colab": { "name": "tpu.ipynb", "toc_visible": true, "machine_shape": "hm", "gpuType": "V28" }, "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 0 }