Tensorflow에서 multi GPU 사용 방법

1. GPU 확인

• 보유하고 있는 gpu 정보와 해당 gpu 번호 확인

from tensorflow.python.client import device_lib
device_lib.list_local_devices()

• 아래 결과를 통해 gpu가 2개 있는 것을 확인할 수 있다.

 

MirroredStrategy

2. 사용 가능한 모든 GPU 병렬처리

• MirroredStrategy 인스턴스 생성 시 아무것도 지정하지 않으면 multi GPU에 대해 분산처리되어 동작함.
• 마찬가지로 with mirrored_strategy.scope() 블럭 내에서는 모델 생성 및 학습 모델 설계 관련 구현하고, with 구문 내에서 model이 compile되어야함.

import tensorflow as tf
mirrored_strategy = tf.distribute.MirroredStrategy()            # 비어있는 상태로 사용
 
def model1():
    inp = layers.Input([10, 30], name="input")
     
    rnn = layers.Bidirectional(layers.LSTM(64, return_sequences=True), merge_mode="concat")(inp)
    rnn = layers.LayerNormalization(epsilon=1e-4)(rnn)
     ...
    out = layers.Dense(9, activation="softmax")(out)
    return models.Model(inputs=[inp], outputs=[out])
 
with mirrored_strategy.scope():
     model = model1()
     model.compile(...)
     ...
 
model.fit(
    dataset,
    validation_data=testset,
    callbacks = [lr_scheduler, earlystopping, checkpoint],
    epochs=1000
)

• 모델 학습 시 gpu 0, 1번에 대해 병렬처리 되고 있는 것을 확인

$ watch nvidia-smi

3. 특정 GPU 지정하여 사용

• 1번 GPU (/GPU:1)로 지정해서 사용
• 가용할 수 있는 GPU가 2개이므로 특정 GPU 지정해서 모델을 학습시키면 GPU를 사용해서 2개의 모델을 동시에 학습시킬 수 있다.
• with mirrored_strategy.scope() 블럭 내에서는 모델 생성 및 학습 모델 설계 관련 구현
  • with 구문 내에서 모델이 compile되어야함!

import tensorflow as tf
mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/GPU:1"])
 
def model1():
    inp = layers.Input([10, 30], name="input")
     
    rnn = layers.Bidirectional(layers.LSTM(64, return_sequences=True), merge_mode="concat")(inp)
    rnn = layers.LayerNormalization(epsilon=1e-4)(rnn)
     ...
    out = layers.Dense(9, activation="softmax")(out)
    return models.Model(inputs=[inp], outputs=[out])
 
with mirrored_strategy.scope():
     model = model1()
     model.compile(...)
     ...
 
model.fit(
    dataset,
    validation_data=testset,
    callbacks = [lr_scheduler, earlystopping, checkpoint],
    epochs=1000
)

• 모델 학습 시 1번 gpu만 사용하고 있는 것을 확인

$ watch nvidia-smi