このTPUトリックでTensorFlowトレーニング速度を50%増やす

コンテンツ概要 セットアップ TPU初期化」 マニュアル装置の配置 配布戦略 TPUの分類 Kerasモデルを定義する データセット Load the Dataset Keras High-Level APIを使用してモデルをトレーニングする トレーニングループを使用してモデルをトレーニング tf.function 内の複数のステップでパフォーマンスを向上させる 次のステップ このガイドは、基本的なトレーニングを実行する方法を示しています。 and TPU Pods, a collection of TPU devices connected by dedicated high-speed network interfaces, with 「Custom Training Loops」 Tensor Processing Units(TPU) tf.keras TPUは、機械学習ワークロードを加速するために使用されるGoogleのカスタム開発アプリケーション特有の統合回路(ASIC)です。 , the そして、 . Google コラボ TPU研究クラウド クラウド TPU セットアップ この Colab ノートブックを実行する前に、ノートブックの設定をチェックすることで、ハードウェア加速器が TPU であることを確認してください。 > > > > > > . Runtime Change runtime type Hardware accelerator TPU TensorFlow Datasets を含むいくつかの必要なライブラリをインポートする: 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は典型的に 従って、ユーザーのPythonプログラムを実行するローカルプロセスとは異なりますので、リモートクラスターに接続してTPUを初期化するために、いくつかの初期化作業を行う必要があります。 論点 2 あなたがGoogle Compute Engine (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による分散型トレーニング 使用する The option implements synchronous distributed training. TPUs provide their own implementation of efficient all-reduce and other collective operations across multiple TPU cores. TPUs provide their own implementation of efficient all-reduce and other collective operations across multiple TPU cores. TPUs provide their own implementation of efficient all-reduce and other collective operations across multiple TPU cores. TPUs provide their own implementation of efficient all-reduce and other collective operations across multiple TPU cores. . tf.distribute.TPUStrategy TPUStrategy これを証明するには、Create a 対象: 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. 以下は、同じ入力を受け取るすべてのコアを示す例です。 そして、それぞれのコアでマトリックスの倍数を独立して実行します. The outputs will be the values from all the replicas. 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 でトレーニングする場合に使用するものとは異なりません。 , so the variables can be created on each TPU device. コードの他の部分は、内部に存在する必要はありません。 スコープ 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)]) データセット Load the Dataset 効率的な利用は、 API は、クラウド TPU を使用する際に重要です. データセットのパフォーマンスについて詳しくは、 . tf.data.Dataset インプットパイプラインパフォーマンスガイド あなたが使用している場合 すべてのデータファイルを TensorFlow で読み取る必要があります。 イン あなたが使用している場合 TPU Nodes および TPU VMs についての詳細については、このページを参照してください。 文書化 TPUノード Dataset Google Cloud Storage(GCS) TPU VMS TPUシステム建築 ほとんどの使用の場合、データをコンバータに変換することをお勧めします。 format and use a 読んでみて(チェック) これを行う方法の詳細については、厳しい要件ではありませんし、他のデータセットリーダーを使用できます。 または . TFRecord tf.data.TFRecordDataset TFRecord and 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 High-Level APIを使用してモデルをトレーニングする あなたはKerasであなたのモデルを訓練することができます そして APIs このステップでは TPU 特有のものは何もありません - 複数の GPU を使用しているかのようにコードを書きます。 代わりに、The . You can learn more in the チュートリアル 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. You can use the API を配布する データセット関数を指定した場合、下の例では、バッチサイズが グローバル パッチ サイズの代わりに per-replica パッチ サイズです。 チュートリアル tf.function tf.distribute Strategy.experimental_distribute_datasets_from_function tf.data.Dataset Dataset Custom トレーニング with 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内の複数のステップを実行することによって これは、巻き込むことによって達成される。 電話 A 内部 , and AutoGraph will convert it to a TPUの職員について. You can learn more about S in the ガイド tf.function Strategy.run tf.range tf.function tf.while_loop tf.function より良いパフォーマンス with tf.function 改善されたパフォーマンスにもかかわらず、この方法は、1 つのステップの内側で実行することに比べて妥協があります。 複数のステップを 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 について、そしてそれらを使用する方法についてもっと知るには: Google Cloud TPU:Google Cloud 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 Notebooks: End-to-end トレーニング サンプル Google Cloud TPU パフォーマンスガイド:アプリケーションの Cloud TPU 構成パラメータを調整することで、Cloud TPU のパフォーマンスをさらに向上させる Distributed training with TensorFlow: How to use distribution strategies — including tf.distribute.TPUStrategy — with examples showing best practices. TensorFlow で分散型トレーニング: ディストリビューション戦略を使用する方法 — including tf.distribute.TPUStrategy — with examples showing best practices. TPU 埋め込み:TensorFlow には、tf.tpu.experimental.embedding 経由で TPU でのトレーニング 埋め込みのための専門的なサポートが含まれています。さらに、TensorFlow Recommenders には tfrs.layers.embedding.TPUEmbedding が含まれています。 TPU Research Cloud(TRC):TRCは、研究者が1000台を超えるクラスターのCloud TPUデバイスへのアクセスを申請することを可能にします。 最初はTensorFlowのウェブサイトに掲載され、この記事はここに新しいタイトルの下に表示され、CC BY 4.0でライセンスを受けています。 最初はTensorFlowのウェブサイトに掲載され、この記事はここに新しいタイトルの下に表示され、CC BY 4.0でライセンスを受けています。 TensorFlowについて TensorFlowについて