在 AI 项目开始的实验阶段的时候，需要做一些实验来验证可行性，通常情况下，实验过程没有被记录起来，导致难以优化和追踪。本文介绍如何通过 mlflow 来记录每次实验的结果、如何注册模型以及如何加载已经注册的模型
实验流程图
┌───────────┐ ┌──────────────┐ ┌─────────────┐ ┌────────────┐
│ Log Model ├─────►│ Regist Model ├─────►│ Change Stage├────►│ Load Model │
└───────────┘ └──────────────┘ └─────────────┘ └────────────┘
Log Model
第一步我们需要做的就是将每次 train 的参数和模型记录到 mlflow 中，相关代码如下：
import mlflow
mlflow.set_tracking_uri("")
mlflow.set_experiment("")
# Auto log parameter and model to mlflow
mlflow.sklearn.autolog()
# 建立模型
model = LinearRegression()
import numpy as np
from sklearn.linear_model import LinearRegression
# 训练数据
X = np.array([[110, 5], [160, 10], [210, 15], [260, 20], [310, 25]])
y = np.array([200, 250, 300, 350, 400])
model.fit(X, y)
每次 train 都会将 model 以一个新的版本记录到 mlflow 中，我们可以用特定的衡量指标，筛选出表现不错的版本，进行模型注册
Regist Model
上一部每次执行，结果都会保存到一个 run 中，每个 run 有一个唯一的 id，通过这个 id，可以对指定的 model 进行注册
import mlflow
mlflow.set_tracking_uri("")
mlflow.set_experiment("")
logged_model = 'runs:/run_id/model'
# Load model as a PyFuncModel.
loaded_model = mlflow.pyfunc.load_model(logged_model)
# 注册模型
registered_model = mlflow.register_model(logged_model, "model_name")
print(f"Model version: {registered_model.version}")
注册的 model 版本号从 1 开始增长，每次注册版本号 +1
Change Stage
可以将某个版本的 model 注册到 stage 或者 production 环境：
# 更改已注册模型的阶段
model_name = "model_name"
version = 2
client = mlflow.MlflowClient()
client.transition_model_version_stage(
name=model_name,
version=version,
stage="Production"
)
Load Model
在 model serving 阶段，可以从 mlflow 加载特定 stage 的 model：
model_name = "model_name"
# model_stage = "Production"
model_stage = "Staging"
model = mlflow.sklearn.load_model(
model_uri=f"models:/{model_name}/{model_stage}"
)
# Predict on a Pandas DataFrame.
import pandas as pd
import numpy as np
X = np.array([[100, 5], [150, 10], [200, 15], [250, 20], [300, 25]])
# y = np.array([200, 250, 300, 350, 400])
predict = model.predict(pd.DataFrame(X))
print("预测结果：", predict)
总结
如上为最简单的应用流程，但是在实际应用中，还需要注意如下问题：
- mlflow 默认没有认证和授权的功能，实际落地的时候，需要另外处理
- mlflow 的数据分为 backend store 和 artifact store，当使用一些大模型的时候，需要考虑将 artifact store 与 beckend store 分开，避免相互影响
