Keras 멀티 GPU 사용 방법

Keras로 멀티 GPU를 사용하려고 실제 훈련을 시켜보면 GPU 한 개만 점유해서 사용하여 멀티 GPU를 활용하지 못하는 경우를 볼 수 있다. 간단한 이유는 Keras와 Tensorflow 버전 호환성 문제이기 때문이다.

이럴때 간혹 외국 사이트나 다른 블로그를 보면 Tensorflow 버전을 1.5, 1.4로 낮추면 된다고 써져있는데 전혀 그럴 필요가 없다. Tensorflow 2.6 이상부터 Keras가 별개 패키지로 관리되기 때문에 이전 버전은 유연하게 대응하도록 하자.

 

환경

• Keras==2.3.1
• tensorflow==2.2.0

 

대응 방법

1. Keras 패키지의 Tensorflow Backend의 _get_available_gpus 함수를 Overwrite 한다.

import keras.backend.tensorflow_backend as tfback

def _get_available_gpus():
    """Get a list of available gpu devices (formatted as strings).

    # Returns
        A list of available GPU devices.
    """
    # global _LOCAL_DEVICES
    if tfback._LOCAL_DEVICES is None:
        devices = tf.config.list_logical_devices()
        tfback._LOCAL_DEVICES = [x.name for x in devices]
    return [x for x in tfback._LOCAL_DEVICES if 'device:gpu' in x.lower()]

2. 그래프 비활성화 (disable_eager_execution) 처리

import tensorflow as tf

tf.compat.v1.disable_eager_execution()

3. 모델에 multi_gpu_model 감싸기

compile 전에 함수로 감싸줘야합니다.

from keras.utils import multi_gpu_model

model = multi_gpu_model(model)

4. optimizers 함수 tensorflow 사용하기

#from keras.optimizers import Adam (X)
from tensorflow.keras.optimizers import Adam

EPOCHS = 100
INIT_LR = 1e-3

opt = Adam(lr=INIT_LR, decay=INIT_LR / EPOCHS)