通过这些TPU技巧提高TensorFlow训练速度50%

内容概览 设置 TPU初始化” 手动装置配置 分配策略 TPU的分类 定义一个Keras模型 加载数据集 使用Keras高级API训练模型 使用自定义训练循环训练模型 通过 tf.function 中的多个步骤提高性能 下一步步骤 本指南展示了如何进行基本培训 TPU Pods 是由专用高速网络接口连接的 TPU 设备的集合。 教练路径。 传感器处理单元(TPU) tf.keras TPU是谷歌定制开发的应用程序特定的集成电路(ASIC),用于加速机器学习工作负载。 , 该 ,和 . 谷歌COLAB TPU研究云 云 TPU 设置 在运行 Colab 笔记本之前,请通过检查笔记本设置,确保您的硬件加速器是 TPU: 」 」 」 . Runtime Change runtime type Hardware accelerator TPU 导入一些必要的库,包括 TensorFlow 数据集: import tensorflow as tf import os import tensorflow_datasets as tfds 2023-06-09 12:13:32.486552: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT TPU初始化 TPU通常是 因此,你需要做一些初始化工作来连接到远程集群并初始化TPU。 论点2 如果您在 Google 计算机 (GCE) 上运行代码,您应该代替将其传输到您的云 TPU 名称中。 云 TPU tpu tf.distribute.cluster_resolver.TPUClusterResolver 注意:TPU初始化代码必须在程序的开始。 TPU初始化代码必须在程序的开始。 Note: resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='') tf.config.experimental_connect_to_cluster(resolver) # This is the TPU initialization code that has to be at the beginning. tf.tpu.experimental.initialize_tpu_system(resolver) print("All devices: ", tf.config.list_logical_devices('TPU')) INFO:tensorflow:Deallocate tpu buffers before initializing tpu system. 2023-06-09 12:13:34.011755: E tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:266] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected INFO:tensorflow:Deallocate tpu buffers before initializing tpu system. INFO:tensorflow:Initializing the TPU system: grpc://10.25.167.66:8470 INFO:tensorflow:Initializing the TPU system: grpc://10.25.167.66:8470 INFO:tensorflow:Finished initializing TPU system. INFO:tensorflow:Finished initializing TPU system. All devices: [LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:0', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:1', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:2', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:3', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:4', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:5', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:6', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:7', device_type='TPU')] 手动装置配置 在 TPU 初始化后,您可以使用手动设备放置来将计算放置在单个 TPU 设备上: a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) with tf.device('/TPU:0'): c = tf.matmul(a, b) print("c device: ", c.device) print(c) c device: /job:worker/replica:0/task:0/device:TPU:0 tf.Tensor( [[22. 28.] [49. 64.]], shape=(2, 2), dtype=float32) 分配策略 通常,您可以在多个 TPU 上以数据并行的方式运行您的模型. 若要在多个 TPU 上(以及多个 GPU 或多个机器上)分发您的模型,TensorFlow 提供 API. 您可以更换您的分布策略,并且该模型将在任何特定(TPU)设备上运行。 指南。 tf.distribute.Strategy 通过TensorFlow进行分布式培训 使用的 选项实现了同步分布式培训。TPU提供自己实现的高效全缩减和其他集体操作跨多个TPU内核,用于 . tf.distribute.TPUStrategy TPUStrategy 为了证明这一点,创建一个 对象: tf.distribute.TPUStrategy strategy = tf.distribute.TPUStrategy(resolver) INFO:tensorflow:Found TPU system: INFO:tensorflow:Found TPU system: INFO:tensorflow:*** Num TPU Cores: 8 INFO:tensorflow:*** Num TPU Cores: 8 INFO:tensorflow:*** Num TPU Workers: 1 INFO:tensorflow:*** Num TPU Workers: 1 INFO:tensorflow:*** Num TPU Cores Per Worker: 8 INFO:tensorflow:*** Num TPU Cores Per Worker: 8 INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0) 要复制计算,以便在所有 TPU 内核中运行,您可以将其传输到 API. 下面是一个示例,显示所有接收相同输入的核心 并且在每个核心上独立执行矩阵倍率,输出将是所有复制品的值。 Strategy.run (a, b) @tf.function def matmul_fn(x, y): z = tf.matmul(x, y) return z z = strategy.run(matmul_fn, args=(a, b)) print(z) PerReplica:{ 0: tf.Tensor( [[22. 28.] [49. 64.]], shape=(2, 2), dtype=float32), 1: tf.Tensor( [[22. 28.] [49. 64.]], shape=(2, 2), dtype=float32), 2: tf.Tensor( [[22. 28.] [49. 64.]], shape=(2, 2), dtype=float32), 3: tf.Tensor( [[22. 28.] [49. 64.]], shape=(2, 2), dtype=float32), 4: tf.Tensor( [[22. 28.] [49. 64.]], shape=(2, 2), dtype=float32), 5: tf.Tensor( [[22. 28.] [49. 64.]], shape=(2, 2), dtype=float32), 6: tf.Tensor( [[22. 28.] [49. 64.]], shape=(2, 2), dtype=float32), 7: tf.Tensor( [[22. 28.] [49. 64.]], shape=(2, 2), dtype=float32) } TPU的分类 涵盖了基本概念后,请考虑一个更具体的例子,本节展示了如何使用分布策略。 在云 TPU 上训练一个 Keras 模型。 tf.distribute.TPUStrategy 定义一个Keras模型 從 A 的定義開始 Keras 模型用于 MNIST 数据集上的图像分类。它与您在 CPU 或 GPU 上进行培训时所使用的模型没有区别。 ,因此可在每个 TPU 设备上创建变量,代码的其他部分不必在 范围。 Sequential Strategy.scope Strategy def create_model(): return tf.keras.Sequential( [tf.keras.layers.Conv2D(256, 3, activation='relu', input_shape=(28, 28, 1)), tf.keras.layers.Conv2D(256, 3, activation='relu'), tf.keras.layers.Flatten(), tf.keras.layers.Dense(256, activation='relu'), tf.keras.layers.Dense(128, activation='relu'), tf.keras.layers.Dense(10)]) 加载数据集 对其有效的使用 API 在使用云 TPU 时至关重要,您可以了解有关数据集性能的更多信息。 . tf.data.Dataset 输入管道性能指南 如果你正在使用 ,您需要存储所有通过 TensorFlow 读取的数据文件 在 如果你正在使用 ,您可以随时随地存储数据. 有关 TPU 节点和 TPU VM 的更多信息,请参阅 文档。 TPU节点 Dataset 谷歌云存储库(GCS) TPU VMs TPU系统架构 对于大多数使用情况,建议将您的数据转换为 格式和使用 a 阅读: 查看 有关如何做到这一点的详细信息,这不是一个艰难的要求,您可以使用其他数据集阅读器,例如: 或 . TFRecord tf.data.TFRecordDataset TFRecord 和 tf.Example 教程 tf.data.FixedLengthRecordDataset tf.data.TextLineDataset 您可以将整个小数据集加载到内存中使用 . tf.data.Dataset.cache 无论使用的数据格式如何,我们强烈建议您使用大文件,大约为 100MB,这在这个网络设置中尤为重要,因为打开文件的重量更高。 如下代码所示,您应该使用 Tensorflow 数据集。 模块,以获取MNIST培训和测试数据的副本。 指定使用可在公共 GCS 桶中使用的副本. 如果您不指定此,TPU 将无法访问下载的数据。 tfds.load try_gcs def get_dataset(batch_size, is_training=True): split = 'train' if is_training else 'test' dataset, info = tfds.load(name='mnist', split=split, with_info=True, as_supervised=True, try_gcs=True) # Normalize the input data. def scale(image, label): image = tf.cast(image, tf.float32) image /= 255.0 return image, label dataset = dataset.map(scale) # Only shuffle and repeat the dataset in training. The advantage of having an # infinite dataset for training is to avoid the potential last partial batch # in each epoch, so that you don't need to think about scaling the gradients # based on the actual batch size. if is_training: dataset = dataset.shuffle(10000) dataset = dataset.repeat() dataset = dataset.batch(batch_size) return dataset 使用Keras高级API训练模型 你可以用Keras训练你的模型 和 此步骤中没有任何 TPU 特定的内容 - 您会写代码,就像您正在使用多个 GPU 和一个 相反,与 你可以学到更多在 教程。 Model.fit Model.compile MirroredStrategy TPUStrategy 与Keras的分布式培训 with strategy.scope(): model = create_model() model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['sparse_categorical_accuracy']) batch_size = 200 steps_per_epoch = 60000 // batch_size validation_steps = 10000 // batch_size train_dataset = get_dataset(batch_size, is_training=True) test_dataset = get_dataset(batch_size, is_training=False) model.fit(train_dataset, epochs=5, steps_per_epoch=steps_per_epoch, validation_data=test_dataset, validation_steps=validation_steps) Epoch 1/5 300/300 [==============================] - 17s 32ms/step - loss: 0.1235 - sparse_categorical_accuracy: 0.9620 - val_loss: 0.0462 - val_sparse_categorical_accuracy: 0.9856 Epoch 2/5 300/300 [==============================] - 7s 24ms/step - loss: 0.0333 - sparse_categorical_accuracy: 0.9894 - val_loss: 0.0401 - val_sparse_categorical_accuracy: 0.9878 Epoch 3/5 300/300 [==============================] - 7s 24ms/step - loss: 0.0186 - sparse_categorical_accuracy: 0.9938 - val_loss: 0.0352 - val_sparse_categorical_accuracy: 0.9900 Epoch 4/5 300/300 [==============================] - 7s 25ms/step - loss: 0.0127 - sparse_categorical_accuracy: 0.9957 - val_loss: 0.0482 - val_sparse_categorical_accuracy: 0.9879 Epoch 5/5 300/300 [==============================] - 7s 24ms/step - loss: 0.0111 - sparse_categorical_accuracy: 0.9962 - val_loss: 0.0448 - val_sparse_categorical_accuracy: 0.9894 为了减少Python的过度使用,并最大限度地提高您的TPU的性能,请通过 争论克里斯 在这个例子中,它增加了约50%的输出量: steps_per_execution Model.compile with strategy.scope(): model = create_model() model.compile(optimizer='adam', # Anything between 2 and `steps_per_epoch` could help here. steps_per_execution = 50, loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['sparse_categorical_accuracy']) model.fit(train_dataset, epochs=5, steps_per_epoch=steps_per_epoch, validation_data=test_dataset, validation_steps=validation_steps) Epoch 1/5 300/300 [==============================] - 14s 45ms/step - loss: 0.1306 - sparse_categorical_accuracy: 0.9591 - val_loss: 0.0420 - val_sparse_categorical_accuracy: 0.9863 Epoch 2/5 300/300 [==============================] - 3s 10ms/step - loss: 0.0333 - sparse_categorical_accuracy: 0.9900 - val_loss: 0.0502 - val_sparse_categorical_accuracy: 0.9846 Epoch 3/5 300/300 [==============================] - 3s 10ms/step - loss: 0.0193 - sparse_categorical_accuracy: 0.9936 - val_loss: 0.0406 - val_sparse_categorical_accuracy: 0.9879 Epoch 4/5 300/300 [==============================] - 3s 10ms/step - loss: 0.0135 - sparse_categorical_accuracy: 0.9955 - val_loss: 0.0416 - val_sparse_categorical_accuracy: 0.9882 Epoch 5/5 300/300 [==============================] - 3s 10ms/step - loss: 0.0110 - sparse_categorical_accuracy: 0.9962 - val_loss: 0.0463 - val_sparse_categorical_accuracy: 0.9882 使用自定义训练循环训练模型 您还可以创建和训练您的模型使用 和 APIs 直接. 您可以使用 火焰分布 数据集函数:请注意,在下面的示例中,批量大小传入 是每复制的批量大小,而不是全球批量大小。 教程。 tf.function tf.distribute Strategy.experimental_distribute_datasets_from_function tf.data.Dataset Dataset 定制培训与 tf.distribute.Strategy 首先,创建模型、数据集和 s : tf.function # Create the model, optimizer and metrics inside the `tf.distribute.Strategy` # scope, so that the variables can be mirrored on each device. with strategy.scope(): model = create_model() optimizer = tf.keras.optimizers.Adam() training_loss = tf.keras.metrics.Mean('training_loss', dtype=tf.float32) training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy( 'training_accuracy', dtype=tf.float32) # Calculate per replica batch size, and distribute the `tf.data.Dataset`s # on each TPU worker. per_replica_batch_size = batch_size // strategy.num_replicas_in_sync train_dataset = strategy.experimental_distribute_datasets_from_function( lambda _: get_dataset(per_replica_batch_size, is_training=True)) @tf.function def train_step(iterator): """The step function for one training step.""" def step_fn(inputs): """The computation to run on each TPU device.""" images, labels = inputs with tf.GradientTape() as tape: logits = model(images, training=True) loss = tf.keras.losses.sparse_categorical_crossentropy( labels, logits, from_logits=True) loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size) grads = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(list(zip(grads, model.trainable_variables))) training_loss.update_state(loss * strategy.num_replicas_in_sync) training_accuracy.update_state(labels, logits) strategy.run(step_fn, args=(next(iterator),)) WARNING:tensorflow:From /tmpfs/tmp/ipykernel_9094/1509474074.py:14: StrategyBase.experimental_distribute_datasets_from_function (from tensorflow.python.distribute.distribute_lib) is deprecated and will be removed in a future version. Instructions for updating: rename to distribute_datasets_from_function WARNING:tensorflow:From /tmpfs/tmp/ipykernel_9094/1509474074.py:14: StrategyBase.experimental_distribute_datasets_from_function (from tensorflow.python.distribute.distribute_lib) is deprecated and will be removed in a future version. Instructions for updating: rename to distribute_datasets_from_function 然后,运行训练循环: steps_per_eval = 10000 // batch_size train_iterator = iter(train_dataset) for epoch in range(5): print('Epoch: {}/5'.format(epoch)) for step in range(steps_per_epoch): train_step(train_iterator) print('Current step: {}, training loss: {}, accuracy: {}%'.format( optimizer.iterations.numpy(), round(float(training_loss.result()), 4), round(float(training_accuracy.result()) * 100, 2))) training_loss.reset_states() training_accuracy.reset_states() Epoch: 0/5 Current step: 300, training loss: 0.1465, accuracy: 95.4% Epoch: 1/5 Current step: 600, training loss: 0.035, accuracy: 98.94% Epoch: 2/5 Current step: 900, training loss: 0.0197, accuracy: 99.39% Epoch: 3/5 Current step: 1200, training loss: 0.0126, accuracy: 99.59% Epoch: 4/5 Current step: 1500, training loss: 0.0109, accuracy: 99.64% 通过内部多个步骤提高性能 tf.function 您可以通过在一个 这可以通过包装 打电话给A 内部 ,并且 AutoGraph 将其转换为 在 TPU 工人. 您可以了解更多关于 S 在 指南。 tf.function Strategy.run tf.range tf.function tf.while_loop tf.function 更好的性能与 tf.function 尽管性能改善,但与在一个单一步骤内运行相比,这种方法存在妥协。 在A中运行多个步骤 它不那么灵活,你不能在步骤中激动地运行东西或任意运行Python代码。 tf.function tf.function @tf.function def train_multiple_steps(iterator, steps): """The step function for one training step.""" def step_fn(inputs): """The computation to run on each TPU device.""" images, labels = inputs with tf.GradientTape() as tape: logits = model(images, training=True) loss = tf.keras.losses.sparse_categorical_crossentropy( labels, logits, from_logits=True) loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size) grads = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(list(zip(grads, model.trainable_variables))) training_loss.update_state(loss * strategy.num_replicas_in_sync) training_accuracy.update_state(labels, logits) for _ in tf.range(steps): strategy.run(step_fn, args=(next(iterator),)) # Convert `steps_per_epoch` to `tf.Tensor` so the `tf.function` won't get # retraced if the value changes. train_multiple_steps(train_iterator, tf.convert_to_tensor(steps_per_epoch)) print('Current step: {}, training loss: {}, accuracy: {}%'.format( optimizer.iterations.numpy(), round(float(training_loss.result()), 4), round(float(training_accuracy.result()) * 100, 2))) Current step: 1800, training loss: 0.009, accuracy: 99.72% 下一步步骤 有关云 TPU 以及如何使用它们的更多信息: 谷歌云TPU:谷歌云TPU主页。 : Google Cloud TPU documentation, which includes: Google Cloud TPU documentation : An overview of working with Cloud TPUs. Introduction to Cloud TPU : Quickstart introductions to working with Cloud TPU VMs using TensorFlow and other main machine learning frameworks. Cloud TPU quickstarts Google Cloud TPU Colab笔记本:终端培训示例。 Google 云 TPU 性能指南:通过调整应用程序的云 TPU 配置参数来进一步提高云 TPU 性能 使用 TensorFlow 的分布式培训:如何使用分布策略 - 包括 tf.distribute.TPUStrategy - 示例显示最佳实践。 TPU 嵌入:TensorFlow 包括通过 tf.tpu.experimental.embedding 在 TPU 上进行培训嵌入的专门支持,此外,TensorFlow Recommenders 包含 tfrs.layers.embedding.TPUEmbedding. 嵌入提供了高效和密集的表示,捕捉了功能之间的复杂性和关系。 TPU 研究云(TRC):TRC 允许研究人员申请访问 1,000 多台云 TPU 设备的集群。 本文最初在 TensorFlow 网站上发表,此文在新标题下出现在这里,并在 CC BY 4.0 版下获得许可。 本文最初在 TensorFlow 网站上发表,此文在新标题下出现在这里,并在 CC BY 4.0 版下获得许可。 腾讯 腾讯