เก็บ Model ของคุณไว้ใช้งานด้วย MLflow

ทำความรู้จัก MLflow

MLflow คือ เครื่องมือที่เอาไว้ช่วย Data Scientist ในการพัฒนา machine learning lifecycle จริงๆมันถูกพัฒนาและใช้งานบน databricks แต่ในที่นี้เราจะพูดถึงตัว open source โดย MLflow มี feature มากมาย ทั้ง model loging, model versioning, model artifact storing หรือแม้กระทั่งการเลือก version code สำหรับ build model ในกรณีที่เรา tune model ด้วยรูปแบบที่แตกต่างกัน ซึ่ง feature ทั้งหมดของ MLflow จะมีดังนี้

1. MLflow Tracking: เป็นตัวเก็บบันทึกผลการทดลองของการ build model แต่ละครั้ง ยกตัวอย่างเช่น ในกรณีที่เราใช้ tensorflow ในการทำ image classification model จะเกิด log ในการ train แต่ละ epoch ว่า metric ต่างๆ ตอนนี้มีค่าเป็นเท่าไร หรือ model parameters ในครั้งนั้นมีค่าอะไรบ้าง เราก็สามารถให้ MLflow Tracking จัดการให้เราได้
2. MLflow Projects: จะเป็นตัวครอบ MLflow Tracking อีกทีหนึ่ง ในกรณีที่เรามีหลาย code version มีหลาย frame work ในการทำ model บางครั้งเราอาจลืมว่าเขียนอะไรไปบ้างหรือมี environment อย่างไร ตัว MLflow Project สามารถจัดการตรงนี้ได้ โดยสามารถเรียก code จาก git repository(code training model ที่เราฝัง MLflow Tracking ไว้แล้ว), set python environment เพื่อมา train ในแต่ละครั้งได้
3. MLflow Models: เมื่อได้ model มาแล้วโดยปกติเราก็ต้อง save model เพื่อนำไปใช้งาน แต่ปัญหาของ Data Scientist คือ มันมี framework เยอะมากสำหรับพัฒนา model เช่น pytorch, tensorflow หรือแม้กระทั่ง scikit-learn โดยแต่ละ framework ก็มีการบันทึกไฟล์ด้วยรูปแบบที่แตกต่างกัน แต่ไม่เป็นไร MLflow Models เกิดมาเพื่อสิ่งนี้ MLflow มีมากกว่า 10 build-in model flavors ให้เลือกใช้ในการจัดเก็บ ไม่เพียงแต่จัดเก็บ artifact model file เพียงอย่างเดียวเท่านั้น MLflow ยังสามารถเพิ่มการเก็บค่าต่างๆในการ train แต่ละครั้ง เช่น ผลการทดลองไม่ว่าจะเป็นรูปแบบ text file หรือ รูปภาพ, เก็บ python environment และ library ที่ใช้, เก็บ architecture ของ model และอื่นๆอีกมากมาย เท่านั้นยังไม่พอ MLflow ยังทำให้การ load model เพื่อมาใช้ง่ายแสนง่ายและเป็นรูปแบบเดียวกัน ที่สำคัญยังสามารถ custom model ที่ไม่มีใน build-in models flavors ได้อีกด้วย
4. MLflow Model Registry: เมื่อมี model ที่ถูกจัดเก็บแล้ว ต่อไปก็ต้องนำมาใช้งาน ในกรณีที่เราเห็นว่าบาง model ของเรามีผลการทดลองที่ดูดีมีอนาคต เราก็ลงทะเบียนให้มันเป็น model เพื่อใช้จริงซะเลย โดย MLflow Model Registry จะเป็นตัว serve model สามารถตั้งสถานะของ model เช่น staging, production (ล่าสุดมีรูปแบบที่เป็น aliases ให้ใช้) และเพิ่ม model version ได้ เพื่อง่ายต่อการนำไปใช้งานจริง
5. ตั้งแต่ version 2.3.2 เริ่มมี MLflow receipt(เป็นตัวทำ ML pipeline) และ version 2.8.1 มี Module ที่ทำงานกับพวก LLM models มากขึ้น อ่านเพิ่มเติม แต่วันนี้เราจะใช้แค่ feature พื้นฐานของ MLflow นั่นคือ Tracking, Models และ Model registry

มาเริ่มใช้งาน MLflow ดีกว่า

Prerequisites:

1. Python environment
2. Install mlflow library(ในตัวอย่างใช้ 2.7.1)

ใน tutorial นี้จะเป็นการแนะนำการใช้งาน MLflow อย่างง่ายเพื่อพัฒนา self-model ของเรา

ภาพที่ 1 MLflow local structure

ขั้นตอนการใช้งาน MLflow

ผมจะพาสร้าง image classification ML model ง่ายๆ โดยจะใช้ประโยชน์จาก feature ของ MLflow หลายๆส่วนร่วมกัน

$ mlflow server --backend-store-uri file:///c:/mlflow/db --default-artifact-root file:///c:/mlflow/artifact --serve-artifacts

step 1 : เปิด command prompt เพื่อสร้าง MLflow tracking server on local สร้าง path สำหรับเก็บ artifact model บน local โดยใช้คำสั่งดังข้างต้น (อย่าปิด cmd จนกว่าจะเลิกใช้)
a) เก็บ backend config ไว้ที่ C:/mlflow/db (ต้องมี directory นี้ก่อนอยู่แล้ว)
b) เก็บ file artifact ไว้ที่ C:/mlflow/storage (ต้องมี directory นี้ก่อนอยู่แล้ว)
c) เมื่อ execute คำสั่งดังกล่าว mlflow จะสร้าง local server ให้เรา copy url นี้ไว้ใช้งานกับ MLflow tracking
d) โดยตัว MLflow tracking ที่เรา setup จะมีการทำงานดังแสดงในภาพที่ 1 ซึ่งเราจะ set root file storage ไว้ที่คอมพิวเตอร์ของเรา

step 2 : ทดสอบเข้า url จาก step 1 c) โดยของผมจะเป็น http://localhost:5000 ซึ่งเป็น website หน้าตาแบบภาพที่ 2

ภาพที่ 2 localhost MLflow

step 3 : สร้าง experiment ขึ้นมาใหม่ตั้งชื่อว่า ml-model ซึ่งสามารถ set artifact directory ได้ แต่ถ้าเราไม่ใส่ MLflow จะ auto generate มาให้ โดยในที่นี้จะเป็น C:/mlflow/artifact/326755177053281989 ตามภาพที่ 3

ภาพที่ 3 create new experiment

step 4 : เราจะสร้าง model เพื่อทำนายรูปของเสื้อผ้าตาม Classify images of clothing แต่เราจะฝัง code เพื่อให้ training job ของเราคุยกับ MLflow และจัดเก็บค่าต่างๆ โดยใช้ tensorflow flavor ของ MLflow Models โดยใช้ code ดังนี้

import tensorflow as tf
import mlflow 
mlflow.set_tracking_uri('http://localhost:5000')
mlflow.set_experiment(experiment_id=326755177053281989)
mlflow.tensorflow.autolog()
fashion_mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10)
])
model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['accuracy'])
model.fit(train_images, train_labels, epochs=10)

1. ใช้คำสั่ง mlflow.set_tracking_uri('http://localhost:5000') เพื่อให้ training job ของเราติดต่อกับ mlflow tracking server ที่เราสร้างขึ้นมา
2. mlflow.set_experiment(326755177053281989)คือคำสั่งให้เราจัดเก็บผลการทดลองไว้ที่ experiment_name = ml-model หรือ experiment_id = 326755177053281989 ที่เราสร้างไว้ใน step 3
3. เราใช้ feature ของ tensorflow flavor โดยให้ training job เก็บค่าต่างที่ใช้ในการ train model ครั้งนี้ โดยใช้คำสั่ง mlflow.tensorflow.autolog()

step 5 : เมื่อเข้าไปที่ http://localhost:5000 จะปรากฎ model ของเรา และจัดเก็บข้อมูลต่างๆที่เกิดจากการทำงานร่วมกันของ MLflow Tracking และ MLflow Models ดังภาพที่ 4

ภาพที่ 4 MLflow ที่จัดเก็บ model ของเรา

step 6 : เมื่อมีผลการทดลองของ training model เราสามารถลงทะเบียน model โดยใช้ feature ของ MLflow Model Registry เพื่อใช้เป็น model serving ซึ่งเราจะตั้งชื่อ model ว่า clothing-classification และตั้ง stage ของ model นี้เป็น Staging ดังภาพที่ 5

ภาพที่ 5 ลงทะเบียน model

step 7 : ขั้นต่อไปคือการนำมาใช้งาน อย่างที่บอก MLflow มันจัดการความแตกต่างของ model flavors ไว้ในรูปแบบเดียวกัน ดังนั้นถ้าจะนำมาใช้ก็เพียงคุยผ่าน mlflow ด้วยคำสั่งต่อไปนี้

import mlflow 
import tensorflow as tf
fashion_mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
 
mlflow.set_tracking_uri('http://localhost:5000')
model_url = 'models:/clothing-classification/Staging'
model = mlflow.tensorflow.load_model(model_url)
prediction = model.predict(test_images)

เพียงแค่ใช้คำสั่ง mlflow.tensorflow.load_model(model_url)ก็สามารถ load model มาใช้งานได้เลย ที่เหลือแค่เพียงใช้ module function เพื่อ predict ตามรูปแบบของแต่ละ frame work

จะเห็นว่าการทำงานส่วนต่างๆกว่าจะได้ model ใดๆขึ้นมา จะใช้ feature การทำงานของ MLflow หลายๆส่วนร่วมกัน แต่ที่สุดแล้วก็ได้ model ที่เก็บในรูปแบบที่เราใช้ค่อนข้างง่าย สามารถเข้าไปดูผลการทดลองหรือส่วนอื่นๆได้ในภายหลัง

ตอนนี้เราก็สามารถใช้ MLflow ในการเก็บ model ไว้ใช้งานในภายหลัง แม้ในตัวอย่างจะเป็น local server แต่ถ้าใช้งานจริงวิธีการก็คล้ายกับขั้นตอนที่ผ่านมาเพียงแค่เปลี่ยน tracking_uri เท่านั้น

ตัวอย่าง code สามารถไปดูได้ที่

mlpipeline-tutorial/mlflow/base-mlflow at main · Janebhop/mlpipeline-tutorial

บทความในซีรีย์นี้
part 1: เก็บ Model ของคุณไว้ใช้งานด้วย MLflow
part 2: MLflow custom model flavor
part 3: Serve MLflow model with BentoML

Acknowledgment
ทีม Data Scientist บริษัท เงินเทอร์โบ จำกัด(มหาชน)
และพรี่ชายผม Thawatchai Chandech