[TF 2.x] TensorFlow 2.0 에서 multi GPU 사용하기 - 텐서플로우 문제 해결

 

*update 2020-10-16 - multi_gpu_model -> tf.distribute.MirroredStrategy

 

필요한건 단 두줄입니다!

 

from tensorflow.keras.utils import multi_gpu_model
parallel_model = multi_gpu_model(model, gpus=2)

 

keras의 함수죠! keras 쓰셨던 분은 익숙하실 합수입니다. model을 컴파일 하기 전에 multi_gpu_model로 변환해주고, parallel_model을 컴파일, 그리고 fit 해주면 됩니다! 참 쉽죠??

 

위의 함수는 더이상 지원되지 않는 것 같습니다.

 

Warning: THIS FUNCTION IS DEPRECATED. It will be removed after 2020-04-01. Instructions for updating: Use tf.distribute.MirroredStrategy instead.

 

따라서 tf.distribute.MirroredStrategy를 살펴봅시다. (TF 2.3 이상의 버전에서 사용하시면 안전합니다.)

 

mirrored_strategy = tf.distribute.MirroredStrategy()
#mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])

with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
  model.compile(loss='mse', optimizer='sgd')
  
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10)
model.fit(dataset, epochs=2)
model.evaluate(dataset)

 

주석으로 처리한 부분은 특정 GPU만 사용할 때에 활용하시면 됩니다. 예를들어 gpu 0번과 1번만 사용하도록 하려면 devices에 ["/gpu:0", "/gpu:1"] 를 입력해주면 됩니다!

그런데 GPU는 두개 다 잡히고 에러 없이 잘 돌아가지만~ nvidia-smi로 살펴보면, 혹시 하나의 GPU만 돌고있는 것을 발견하실 수도 있습니다! 그러한 경우에는 아래와 같이 두 줄을 추가해주면 됩니다 :)

 

import tensorflow as tf 
tf.compat.v1.disable_eager_execution()

 

tf2.0에서는 eager mode가 default로 되어있는데, multi GPU를 위한 분산 strategy를 위해서는 disable해주어야합니다!

 

eager mode가 뭔지 궁금하신 분들은, eager execution (즉시 실행)이 설명되어있는 다음 포스트를 참고해주세요!

 

TensorFlow 2.0의 주요 기능 (TensorFlow와 Keras의 장점의 결합)

TensorFlow 2.0의 주요 기능 (TensorFlow와 Keras의 장점의 결합)