使用 mlflow 完成 model tracking、registry

在 AI 项目开始的实验阶段的时候，需要做一些实验来验证可行性，通常情况下，实验过程没有被记录起来，导致难以优化和追踪。本文介绍如何通过 mlflow 来记录每次实验的结果、如何注册模型以及如何加载已经注册的模型

实验流程图

https://asciiflow.com/#/share/eJyrVspLzE1VssorzcnRUcpJrEwtUrJSqo5RqohRsrK0MNCJUaoEsowsLICsktSKEiAnRunRlD3EIAUwIFIxJbpI1BETkwfS5JOfruCbn5Kag03rtF0gJUGp6ZnFJYRUOWck5qWnKgSXJKanYlfhk5%2BYgjAFYv%2BQDkGlWqVaAN0bXmc%3D)

┌───────────┐      ┌──────────────┐      ┌─────────────┐     ┌────────────┐
│ 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)

总结

如上为最简单的应用流程，但是在实际应用中，还需要注意如下问题：

1. mlflow 默认没有认证和授权的功能，实际落地的时候，需要另外处理
2. mlflow 的数据分为 backend store 和 artifact store，当使用一些大模型的时候，需要考虑将 artifact store 与 beckend store 分开，避免相互影响