複数のGPUを使用してTensorflowを学習する方法についてまとめます。
1. Mirrored Strategy
Mirrored Strategyとは、Tensor Flowで作成したモデルを複数のGPU、TPUを使用して学習するためのAPIです。Mirrored Strategyの公式ページに細かく解説が載っていますが、本ページでは実装までに必要な最小限のコードをまとめます。
2. 実装
インポートはtensorflowのみでOK。
import tensorflow as tf
Mirrored Strategyの宣言とモデルの作成・コンパイル。モデルは適当です。
mirrored_strategy = tf.distribute.MirroredStrategy() with mirrored_strategy.scope(): model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))]) model.compile(loss='mse', optimizer='sgd')
あとは、いつも通りmodel.fitで学習するだけです。model.fitはMirrored Strategyに入れないように気を付けてください。なお、上記の例はすべてのGPUを使用する場合のコードですが、使用するGPUを指定したい場合は以下のようにする必要があります。
mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1", "/gpu:2"]) with mirrored_strategy.scope(): model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))]) model.compile(loss='mse', optimizer='sgd')
