MLFLOW 도전기

MLflow는 End to End로 데이터브릭스에서 만든 머신러닝 라이프 사이클을 관리할 수 있는 오픈소스입니다.
MLOps 관련 오픈소스중에 빠르게 성장 중이라고 하여 Fashion Mnist 데이터셋을 이용해 간단하게 예측 한 번 해보았습니다.
TensorFlow Keras 모델을 학습하고, MLFlow를 사용하여 실험을 기록하도록 하는 예제를 구현해볼께요!

0. 설치 및 환경

우선 주피터를 통해 mlflow를 설치해줍니다.

!pip install mlflow

##________import _________##
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.callbacks import ModelCheckpoint
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.datasets import fashion_mnist
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import sklearn
from sklearn.metrics import roc_auc_score
import mlflow
import mlflow.tensorflow
##_______version_________##
print(f"Tensorflow : {tf.__version__}")
print(f"scikit-learn : {sklearn.__version__}")
print(f"Numpy : {np.__version__}")
print(f"MLFlow : {mlflow.__version__}")
print(f"Matplotlib : {matplotlib.__version__}")

1. 학습을 위한 데이터 셋 확인

mlflow를 설치하였으니 실제로 한번 운영해봐야겠죠?
fashion mnist 데이터셋을 이용해보겠습니다.

(x_train, y_train), (x_test, y_test) = fashion_mnist.load_data()
x_train = x_train.reshape(x_train.shape[0],x_train.shape[1], x_train.shape[2], 1)
x_test = x_test.reshape(x_test.shape[0], x_test.shape[1], x_test.shape[2], 1)

y_train = tf.keras.utils.to_categorical(y_train)
y_test = tf.keras.utils.to_categorical(y_test)

plt.imshow(x_train[0], cmap = 'gray'), print("Class: ", y_train[0])

fashion mnist 데이터에서 첫 번째 이미지를 가져와 확인 해보겠습니다.

훈련 데이터셋과 테스트 데이터셋의 데이터 양과 배열도 확인해볼까요?

print("Shapes")
print(f"x_train : {x_train.shape}")
print(f"y_train : {y_train.shape}")
print(f"x_test : {x_test.shape}")
print(f"y_test : {y_test.shape}")

각 배열의 형태 (shape)를 출력해보면, train은 60000개, test는 10000개 이미지를 가지고 있네요..!
x_train와 y_train은 각각 (28x28)의 형태를 갖고 있고, x_test와 y_test은 각각 0에서 9 사이의 정수를 나타내는 배열 임을 알 수 있습니다.

2. 모델링 및 학습

데이터셋을 확인해봤으니, 모델을 정의 해보겠습니다.

# 모델 정의
model = Sequential([
    tf.keras.layers.Reshape((28, 28, 1), input_shape=(28, 28)),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax'),
])
model.summary()

저는 신경망층을 여러 층을 구성했는데, summary() 명령어를 통해 간단하게 확인하고 컴파일을 이제 해보겠습니다.

# 컴파일
model.compile(optimizer='Adam', loss='categorical_crossentropy', metrics=['acc'])

# 체크포인트
checkpoint_path = 'my_checkpoint.ckpt'
checkpoint = ModelCheckpoint(filepath = checkpoint_path,
                             save_weights_only = True,
                             save_best_only = True,
                             monitor='val_loss',
                             verbose=1)

옵티마이저는 ADAM, 손실함수는 categorical_crossentropy로 측정값은 정확도로 컴파일 했습니다. 체크포인트는 검증셋의 손실값을 가지고 모니터링을 하도록 설정했습니다.

그럼 이 글의 목적이었던 MLFlow를 통해 학습을 진행해보겠습니다.

# 학습 with MLFlow
mlflow.set_experiment("TF_Keras_fashionMnist")

with mlflow.start_run():

    mlflow.tensorflow.autolog()

    model.fit(x_train, y_train,
              validation_data = (x_test, y_test),
              epochs = 20,
              callbacks = [checkpoint],
              )

    model.load_weights(checkpoint_path)

    preds = model.predict(x_test)
    preds = np.round(preds)

    eval_acc = model.evaluate(x_test, y_test)[1]
    auc_score = roc_auc_score(y_test, preds, multi_class='raise')

    print("eval_acc: ", eval_acc)
    print("auc_score: ", auc_score)

    mlflow.tensorflow.mlflow.log_metric('eval_acc', eval_acc)
    mlflow.tensorflow.mlflow.log_metric('auc_score', auc_score)

mlflow.end_run()

mlflow.start_run() 함수를 호출하여 새로운 MLFlow 실험을 시작하고, mlflow.tensorflow.autolog() 함수를 호출하여 TensorFlow 모델의 학습과정을 자동으로 MLFlow로 기록합니다.

모델을 학습하면서 model.load_weights() 함수를 사용하여 이전에 저장된 모델 가중치를 불러옵니다. 이후 모델의 예측값을 계산하고, np.round() 함수를 사용하여 반올림합니다.

마지막으로 모델을 평가하고, MLFlow를 사용하여 eval_acc와 auc_score를 기록하고, mlflow.end_run() 함수를 호출하여 MLFlow를 종료합니다.

3. MLFlow UI를 통한 결과 확인

MLFlow UI에 들어가 해당 결과 값을 확인해보겠습니다.
제가 이름을 정한 TF_Keras_fashionMnist 프로젝트에 들어가면 표를 통해이렇게 결과가 나오네요.

!mlflow ui -p 9999

실행이 성공적으로 기록됐음을 알 수 있습니다.

제가 했던 파라미터와 지표들을 모두 확인할 수 있었는데, 지표를 확인해보니 정확도는 0.956으로 좋은 성적을 기록했네요!

아래와 같은 명령어를 통해 MLFlow에 모델을 적제할 수 있습니다.

#MLFlow 모델 적재
loaded_model = mlflow.keras.load_model("runs:/c0f3fe7f53d74d0aae93b88abc9c45c5/model")

그리고 이전에 기록한 것과 동일한 평가지표로 신속하게 계산할 수 있습니다. 제가 아까 학습했던 모델을 불러와 아래 코드처럼 작성하게 된다면 제가 설정했던 지표, 파라미터로 학습이 진행 되는 것을 확인할 수 있습니다.

eval_loss, eval_acc = loaded_model.evaluate(x_test, y_test)

preds = loaded_model.predict(x_test)
preds = np.round(preds)

eval_auc = roc_auc_score(y_test, preds)

print("Eval Loss: ", eval_loss)
print("Eval Acc: ", eval_acc)
print("Eval Auc: ", eval_auc)

이렇게 Fashion Mnist 데이터셋을 가지고 MLflow를 간편하게 구현해보았습니다.
또, 재밌는 것이 있다면 블로그에 남기도록 할게요!