MLflow 怎样记录模型签名和输入示例

如果模型已经训练完成，但调用方不知道该传哪些列、数据类型是什么，部署后的问题通常会从“模型不准”变成“输入根本没有按约定到达”。MLflow 的做法是把模型签名当成一份输入输出契约：记录模型时传入 input_example，MLflow 会根据这个真实样例推断并保存 signature，同时把样例作为模型工件的一部分保存下来。

官方地址：https://mlflow.org/docs/latest/ml/model/signatures/

先把 input_example 和 signature 分清楚

input_example 回答的是“给模型什么样的一小份真实数据”。例如分类模型需要四列浮点特征，就可以提供包含这四列的一行 DataFrame。它既帮助读者理解输入格式，也让 MLflow 在记录模型时有机会调用模型做输入检查并推断 schema。

signature 回答的是“模型接受什么、返回什么，以及参数是什么”。它可以包含输入 schema、输出 schema 和参数 schema。对于 DataFrame，输入通常表现为列名与类型；对于 NumPy 等张量输入，则更关注 dtype 和 shape。两者配合后，模型目录里既有可读的示例，也有可交给部署工具使用的契约。

用一个最小分类项目准备输入样例

下面的示例选择 Iris 数据集，只为了把“训练特征”和“记录样例”放在同一条清晰链路里。先安装依赖，再启动本地 MLflow Tracking，代码本身不依赖远程服务。

# 安装模型训练、表格处理和本地记录所需依赖
python -m pip install mlflow scikit-learn pandas

# 启动本地 MLflow UI，默认读取当前目录下的 mlruns
mlflow ui --host 127.0.0.1 --port 5000

准备样例时不要随手从原始数组截取一行后再丢掉列名。列名是表格模型签名的重要部分，调用方后续需要按相同名称组织请求。

在 log_model 中一次记录签名和输入示例

关键动作只有一个：在对应 flavor 的 log_model 调用里传入 input_example。下面的代码将一行四列 DataFrame 作为样例，MLflow 会据此推断输入结构；模型的输出则由当前模型预测结果决定。

import mlflow
import mlflow.sklearn
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

# 读取带列名的数据，避免后面的签名丢失特征名称
iris = load_iris(as_frame=True)
X = iris.data
y = iris.target

# 训练一个最小分类模型，示例重点是记录契约而不是调参
model = RandomForestClassifier(n_estimators=20, random_state=7)
model.fit(X, y)

# 取一行真实特征作为可复用的输入示例
input_example = X.iloc[[0]]

mlflow.set_experiment("signature-demo")
with mlflow.start_run() as run:
    # input_example 让 MLflow 推断并保存模型 signature
    mlflow.sklearn.log_model(
        sk_model=model,
        name="iris_classifier",
        input_example=input_example,
    )

    # 记录少量训练信息，便于在 UI 中区分不同实验运行
    mlflow.log_param("feature_count", X.shape[1])
    mlflow.log_param("estimator", "RandomForestClassifier")

    print(f"run_id={run.info.run_id}")

这段代码的重点不是 RandomForestClassifier，而是样例与训练输入使用同一套列名和数据类型。若模型实际上接收经过预处理的张量，样例也要体现预处理后的形状，不能拿原始表格冒充最终接口。

从模型元数据读取签名和样例

记录成功后，可以加载刚才保存的模型并查看元数据。这里的读取动作用于理解契约，不需要重新训练模型。生产代码还可以把这些字段导出给 API schema 或请求校验层，但不要把模型签名和业务字段说明混成一份无法维护的大文档。

import mlflow

# 用实际的模型 URI 加载已记录的 PyFunc 模型
model_uri = "runs://iris_classifier"
loaded_model = mlflow.pyfunc.load_model(model_uri)

# signature 包含输入、输出和参数 schema
signature = loaded_model.metadata.signature
print(signature)

# metadata 还会暴露保存时的输入样例信息
print(loaded_model.metadata.get_input_schema())
print(loaded_model.metadata.get_output_schema())

在输出中重点看四类信息：输入列是否仍叫 sepal length (cm) 等原始列名、列类型是否符合训练时的浮点类型、输出是否是预期的预测结构，以及是否存在可选或必填参数。签名的意义是让这些约定靠近模型工件，而不是只留在聊天记录里。

样例不足时显式构造 ModelSignature

自动推断适合大多数表格模型，但有些输入示例只覆盖了一个分支，或者模型输入是固定 shape 的数组，这时可以显式给出 ModelSignature。显式签名并不会替代输入样例：如果调用方需要一个可以复制的请求范例，仍然建议同时保留 input_example。

import numpy as np
from mlflow.models import ModelSignature
from mlflow.types import Schema, ColSpec, TensorSpec

# 表格模型用列名和类型描述输入输出
tabular_signature = ModelSignature(
    inputs=Schema([
        ColSpec("double", "sepal length (cm)"),  # 训练特征的列名与类型
        ColSpec("double", "sepal width (cm)"),
        ColSpec("double", "petal length (cm)"),
        ColSpec("double", "petal width (cm)"),
    ]),
    outputs=Schema([ColSpec("long")]),  # 分类标签是整数预测结果
)

# 张量模型应把 dtype 与形状一起写进签名
image_signature = ModelSignature(
    inputs=Schema([
        TensorSpec(np.dtype(np.float32), (-1, 28, 28, 1)),  # -1 表示批大小
    ])
)

print(tabular_signature)
print(image_signature)

显式签名要和真正传给模型的对象保持一致。把 DataFrame 的列式接口写成 Tensor，或把批大小固定成一个值，都会让部署端得到错误的契约。若使用较新的 MLflow 与 Python 类型提示，也可以让类型提示参与签名推断，但复杂嵌套输入仍应先设计清楚批量接口。

按错误类型排查签名不匹配

当服务请求被拒绝时，先对照签名，不要立即怀疑模型权重。最常见的问题可以按下面的顺序定位：

一个实用原则是：签名描述“模型真正接收的对象”，不是描述前端表单原样提交的对象。如果中间有特征工程、图片归一化或文本分词，应明确签名落在哪一层，并让示例、模型和部署入口使用同一层级。

记录完成后的发布前清单

1. 模型能被 MLflow 按记录的 URI 加载，且元数据中存在 signature。
2. 输入 schema 的列名、数据类型或 Tensor shape 与实际推理入口一致。
3. input_example 能代表调用方准备的数据，而不是只对训练集内部对象有效。
4. 输出 schema 能被下游服务解释，预测标签、概率或结构化对象的含义没有混淆。
5. 如果模型要进入统一注册或部署平台，提前确认该平台是否要求显式签名。

常见问题

只写 input_example，不手动传 signature 可以吗？

可以。对常见 flavor，MLflow 会根据输入示例推断签名，并将两者保存到模型中。只要示例覆盖了真实输入边界，自动推断通常比手写一份与代码脱节的 schema 更省维护。

一行样例会不会让签名不完整？

有可能。样例主要提供形状、列名和类型；如果某些分支字段只在特殊请求中出现，应显式设计可选字段或补充 ModelSignature，不能期待一行样例自动表达所有业务语义。

DataFrame 和 Tensor 签名最大的区别是什么？

DataFrame 更关注列名与列类型，适合传统表格模型；Tensor 更关注 dtype 和 shape，适合图像、音频或 embedding。选择签名形式时，应以模型真正接收的对象为准。

为什么签名正确，线上仍然报错？

签名只描述模型输入输出契约，不会替你修复请求编码、预处理顺序、模型 URI 或服务端依赖。先确认请求到达模型前没有丢列、改类型或改变 shape，再继续看运行环境问题。