跳到主要内容

MLflow 数据集跟踪教程

mlflow.data 模块是 MLflow 生态系统的重要组成部分，旨在增强您的机器学习工作流程。该模块使您能够在模型训练和评估期间记录和检索数据集信息，利用 MLflow 的跟踪功能。

关键接口​

与 mlflow.data 模块相关的两个主要抽象组件是 Dataset 和 DatasetSource。

Dataset​

Dataset 抽象是一个元数据跟踪对象，它保存了有关给定已记录数据集的信息。

存储在 Dataset 对象中的信息包括特征、目标和预测，以及数据集名称、摘要（哈希）、模式和概要等元数据。您可以使用 mlflow.log_input() API 记录此元数据。该模块提供了从各种数据类型构建 mlflow.data.dataset.Dataset 对象的功能。

此抽象类有许多具体的实现，包括：

以下示例演示了如何从 Pandas DataFrame 构建一个 mlflow.data.pandas_dataset.PandasDataset 对象：

import mlflow.data
import pandas as pd
from mlflow.data.pandas_dataset import PandasDataset


dataset_source_url = "https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/winequality-white.csv"
raw_data = pd.read_csv(dataset_source_url, delimiter=";")

# Create an instance of a PandasDataset
dataset = mlflow.data.from_pandas(
    raw_data, source=dataset_source_url, name="wine quality - white", targets="quality"
)

DatasetSource​

DatasetSource 是给定 Dataset 对象的一个组件，提供了与原始数据源的链接沿袭。

Dataset 的 DatasetSource 组件表示数据集的来源，例如 S3 中的目录、Delta Table 或 URL。它在 Dataset 中被引用，用于理解数据的来源。已记录数据集的 DatasetSource 可以通过访问 Dataset 对象的 source 属性或使用 mlflow.data.get_source() API 来检索。

提示

MLflow 中许多支持自动日志记录 (autologging) 的 flavor 会在记录数据集时自动记录数据集的来源。

注意

下面显示的示例仅用于教学目的，因为在训练运行之外记录数据集并非普遍做法。

示例用法​

以下示例演示了如何使用 log_inputs API 来记录训练数据集、检索其信息并获取数据源：

import mlflow
import pandas as pd
from mlflow.data.pandas_dataset import PandasDataset


dataset_source_url = "https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/winequality-white.csv"
raw_data = pd.read_csv(dataset_source_url, delimiter=";")

# Create an instance of a PandasDataset
dataset = mlflow.data.from_pandas(
    raw_data, source=dataset_source_url, name="wine quality - white", targets="quality"
)

# Log the Dataset to an MLflow run by using the `log_input` API
with mlflow.start_run() as run:
    mlflow.log_input(dataset, context="training")

# Retrieve the run information
logged_run = mlflow.get_run(run.info.run_id)

# Retrieve the Dataset object
logged_dataset = logged_run.inputs.dataset_inputs[0].dataset

# View some of the recorded Dataset information
print(f"Dataset name: {logged_dataset.name}")
print(f"Dataset digest: {logged_dataset.digest}")
print(f"Dataset profile: {logged_dataset.profile}")
print(f"Dataset schema: {logged_dataset.schema}")

上述代码片段的 stdout 结果如下：

Dataset name: wine quality - white
Dataset digest: 2a1e42c4
Dataset profile: {"num_rows": 4898, "num_elements": 58776}
Dataset schema: {"mlflow_colspec": [
    {"type": "double", "name": "fixed acidity"},
    {"type": "double", "name": "volatile acidity"},
    {"type": "double", "name": "citric acid"},
    {"type": "double", "name": "residual sugar"},
    {"type": "double", "name": "chlorides"},
    {"type": "double", "name": "free sulfur dioxide"},
    {"type": "double", "name": "total sulfur dioxide"},
    {"type": "double", "name": "density"},
    {"type": "double", "name": "pH"},
    {"type": "double", "name": "sulphates"},
    {"type": "double", "name": "alcohol"},
    {"type": "long", "name": "quality"}
    ]}

我们也可以导航到 MLflow UI 查看已记录 Dataset 的样子。

当我们想从存储位置加载数据集（调用 load 会在本地下载数据）时，我们可以通过以下 API 访问 Dataset 的来源：

# Loading the dataset's source
dataset_source = mlflow.data.get_source(logged_dataset)

local_dataset = dataset_source.load()

print(f"The local file where the data has been downloaded to: {local_dataset}")

# Load the data again
loaded_data = pd.read_csv(local_dataset, delimiter=";")

上面的 print 语句解析为调用 load 时创建的本地文件。

The local file where the data has been downloaded to:
/var/folders/cd/n8n0rm2x53l_s0xv_j_xklb00000gp/T/tmpuxwtrul1/winequality-white.csv

将 Datasets 与其他 MLflow 功能结合使用​

mlflow.data 模块在将数据集与 MLflow 运行关联方面起着至关重要的作用。除了将 MLflow 运行记录与训练期间使用的数据集关联起来的显而易见的用处之外，MLflow 中还有一些集成允许直接使用已通过 mlflow.log_input() API 记录的 Dataset。

如何将 Dataset 与 MLflow evaluate 结合使用​

注意

Datasets 与 MLflow evaluate 的集成是在 MLflow 2.8.0 中引入的。早期版本不具备此功能。

要了解此集成如何工作，让我们看一个相当简单典型的分类任务。

import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
import xgboost

import mlflow
from mlflow.data.pandas_dataset import PandasDataset


dataset_source_url = "https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/winequality-white.csv"
raw_data = pd.read_csv(dataset_source_url, delimiter=";")

# Extract the features and target data separately
y = raw_data["quality"]
X = raw_data.drop("quality", axis=1)

# Split the data into training and test sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.33, random_state=17
)

# Create a label encoder object
le = LabelEncoder()

# Fit and transform the target variable
y_train_encoded = le.fit_transform(y_train)
y_test_encoded = le.transform(y_test)

# Fit an XGBoost binary classifier on the training data split
model = xgboost.XGBClassifier().fit(X_train, y_train_encoded)

# Build the Evaluation Dataset from the test set
y_test_pred = model.predict(X=X_test)

eval_data = X_test
eval_data["label"] = y_test

# Assign the decoded predictions to the Evaluation Dataset
eval_data["predictions"] = le.inverse_transform(y_test_pred)

# Create the PandasDataset for use in mlflow evaluate
pd_dataset = mlflow.data.from_pandas(
    eval_data, predictions="predictions", targets="label"
)

mlflow.set_experiment("White Wine Quality")

# Log the Dataset, model, and execute an evaluation run using the configured Dataset
with mlflow.start_run() as run:
    mlflow.log_input(pd_dataset, context="training")

    mlflow.xgboost.log_model(
        artifact_path="white-wine-xgb", xgb_model=model, input_example=X_test
    )

    result = mlflow.evaluate(data=pd_dataset, predictions=None, model_type="classifier")
注意

使用 mlflow.evaluate() API 将自动把用于评估的数据集记录到 MLflow 运行中。不需要显式调用来记录输入。

导航到 MLflow UI，我们可以看到数据集、模型、指标以及分类特有的混淆矩阵如何都记录到了运行中。