🌏
Tensorflowで複数GPUを使用する方法
はじめに
本記事ではTensorflowで複数GPUを用いた計算を行うためにTensorflow2.0から使用可能となったtf.distribute.Strategyを用いた例を紹介します。
TensorflowオフィシャルにもGPUの割り当てに関した記事がありますが、こちらはの記事ではより簡潔に解説し、すぐ実装できるという趣旨で紹介していきます。
また、実際に利用される例としてDeep learningがほとんどだと思いますのでDeep learningの学習を例に解説していきます。
Tensorflow Strategy
Tensorflow-gpuを用いたDeep learningの学習ではデフォルトでは単一のGPUが使われます。
マシンに複数のGPUを積んでいて、複数のGPUを用いて学習を行いたい場合、最も簡単な方法はTensorflow2.0から公開されたStrategyを使用します。
また、Strategyを用いた複数GPUの利用はTensorflow公式により推奨されているベストプラクティスです。
使用例
使用方法は非常に簡単でモデルのCompile、BuildをMirroredStrategyのスコープで括ってあげます。
これで複数GPUが適応されます。
今回使用したMirroredStrategy以外にもStrategyが利用でき、複数マシンやTPUへの分散もサポートされています。
ちなみに、Strategyの裏側では、各GPUでモデルのコピーを実行し、それらの間で入力データを分割することで。データの並列処理を行なっています。
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
inputs = tf.keras.layers.Input(shape=(1,))
predictions = tf.keras.layers.Dense(1)(inputs)
model = tf.keras.models.Model(inputs=inputs, outputs=predictions)
model.compile(loss='mse',
optimizer=tf.keras.optimizers.SGD(learning_rate=0.2))
model.fit( ... )
終わり
マルチGPUを用いてDNNに高速に学習させましょうー!
参考
Discussion