Tensorflowで複数GPUを使用する方法

1 min read読了の目安(約1200字

はじめに

本記事では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への分散もサポートされています。

https://www.tensorflow.org/guide/distributed_training?hl=ja

ちなみに、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に高速に学習させましょうー!

参考

https://www.tensorflow.org/guide/gpu?hl=ja
https://www.tensorflow.org/guide/distributed_training?hl=ja