以编程方式从 MlFlow 模型注册表获取模型的输入架构

Question

有没有办法从 MlFlow 模型注册表中获取输入模式（完成训练的特征）？在记录经过训练的模型时，使用“signature”参数捕获输入模式。

Answer 1

我将描述执行此操作的两种方法。

模型签名可以从关联的运行元数据中检索。这是一张显示如何在 UI 中执行此操作的图片：

现在，要以编程方式提取此内容，请注意，记录的模型元数据是在

mlflow.log-model.history

标签下跟踪的。一旦我们知道了相应的运行 ID（我们将其放在手边或查询模型存储），我们就可以使用此代码片段：

import json
import mlflow
from mlflow.client import MlflowClient

client = MlflowClient('http://0.0.0.0:5000')
run_id = '467677aff0074955a4e75492085d52f9'
run = client.get_run(run_id)
log_model_meta = json.loads(run.data.tags['mlflow.log-model.history'])
log_model_meta[0]['signature']

与图一致:-)

{'inputs': '[{"type": "tensor", "tensor-spec": {"dtype": "float64", "shape": [-1, 4]}}]',
 'outputs': '[{"type": "tensor", "tensor-spec": {"dtype": "int64", "shape": [-1]}}]'}

另一种方法是查询模型存储。架构/签名显示在模型视图下，如下所示

可以通过函数

mlflow.models.get_model_info

获取数据，就像在这个片段中

model_uri = client.get_model_version_download_uri('toy-model','10')
model_info = mlflow.models.get_model_info(model_uri)
model_info._signature_dict

Answer 2

我发现 mlflow.types.schema.Schema 也很有帮助！考虑到我想专门访问列名称。

import mlflow
from mlflow.types.schema import Schema

model_uri = f"runs:/{model_run_id}/{run_name}"

# Access the schema from model metadata
schema = model.metadata.signature
feature_names = Schema.input_names(schema.inputs)
target_name = Schema.input_names(schema.outputs)

print(feature_names, target_name)