TensorFlow로 멀티 GPU 분산 트레이닝하기

import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import tensorflow as tf
import keras

# MirroredStrategy 생성
strategy = tf.distribute.MirroredStrategy()
print('Number of devices: {}'.format(strategy.num_replicas_in_sync))

# strategy scope 열기
with strategy.scope():
    # 변수를 생성하는 모든 작업은 strategy scope 내에서 이루어져야 합니다.
    # 일반적으로 모델 생성 및 `compile()`입니다.
    model = Model(...)
    model.compile(...)

    # 사용 가능한 모든 디바이스에서 모델 트레이닝
    model.fit(train_dataset, validation_data=val_dataset, ...)

    # 사용 가능한 모든 디바이스에서 모델 평가
    model.evaluate(test_dataset)

def get_compiled_model():
    # 간단한 2 레이어 Dense 신경망을 만듭니다.
    inputs = keras.Input(shape=(784,))
    x = keras.layers.Dense(256, activation="relu")(inputs)
    x = keras.layers.Dense(256, activation="relu")(x)
    outputs = keras.layers.Dense(10)(x)
    model = keras.Model(inputs, outputs)
    model.compile(
        optimizer=keras.optimizers.Adam(),
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=[keras.metrics.SparseCategoricalAccuracy()],
    )
    return model


def get_dataset():
    batch_size = 32
    num_val_samples = 10000

    # [`tf.data.Dataset`](https://www.tensorflow.org/api_docs/python/tf/data/Dataset) 형식으로 MNIST 데이터셋을 반환합니다.
    (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

    # 데이터를 전처리합니다. (이들은 Numpy 배열입니다)
    x_train = x_train.reshape(-1, 784).astype("float32") / 255
    x_test = x_test.reshape(-1, 784).astype("float32") / 255
    y_train = y_train.astype("float32")
    y_test = y_test.astype("float32")

    # num_val_samples 샘플을 검증용으로 예약합니다.
    x_val = x_train[-num_val_samples:]
    y_val = y_train[-num_val_samples:]
    x_train = x_train[:-num_val_samples]
    y_train = y_train[:-num_val_samples]
    return (
        tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size),
        tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(batch_size),
        tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size),
    )


# MirroredStrategy를 생성합니다.
strategy = tf.distribute.MirroredStrategy()
print("Number of devices: {}".format(strategy.num_replicas_in_sync))

# strategy scope를 엽니다.
with strategy.scope():
    # 변수를 생성하는 모든 작업은 strategy scope 내에서 이루어져야 합니다.
    # 일반적으로 모델 생성 및 `compile()`입니다.
    model = get_compiled_model()

    # 사용 가능한 모든 디바이스에서 모델을 트레이닝합니다.
    train_dataset, val_dataset, test_dataset = get_dataset()
    model.fit(train_dataset, epochs=2, validation_data=val_dataset)

    # 사용 가능한 모든 디바이스에서 모델을 테스트합니다.
    model.evaluate(test_dataset)

# 체크포인트를 저장할 디렉토리를 준비합니다.
checkpoint_dir = "./ckpt"
if not os.path.exists(checkpoint_dir):
    os.makedirs(checkpoint_dir)


def make_or_restore_model():
    # 체크포인트가 있으면 최신 모델을 복원하고,
    # 체크포인트가 없으면 새로운 모델을 생성합니다.
    checkpoints = [checkpoint_dir + "/" + name for name in os.listdir(checkpoint_dir)]
    if checkpoints:
        latest_checkpoint = max(checkpoints, key=os.path.getctime)
        print("Restoring from", latest_checkpoint)
        return keras.models.load_model(latest_checkpoint)
    print("Creating a new model")
    return get_compiled_model()


def run_training(epochs=1):
    # MirroredStrategy를 생성합니다.
    strategy = tf.distribute.MirroredStrategy()

    # strategy scope를 열고 모델을 생성하거나 복원합니다.
    with strategy.scope():
        model = make_or_restore_model()

        callbacks = [
            # 이 콜백은 매 에포크마다 SavedModel을 저장합니다.
            # 현재 에포크를 폴더 이름에 포함시킵니다.
            keras.callbacks.ModelCheckpoint(
                filepath=checkpoint_dir + "/ckpt-{epoch}.keras",
                save_freq="epoch",
            )
        ]
        model.fit(
            train_dataset,
            epochs=epochs,
            callbacks=callbacks,
            validation_data=val_dataset,
            verbose=2,
        )


# 처음 실행 시 모델을 생성합니다.
run_training(epochs=1)

# 같은 함수를 다시 호출하면 이전 상태에서 재개합니다.
run_training(epochs=1)