跳到主要内容

MLflow 数据集跟踪

mlflow.data 模块是贯穿整个机器学习生命周期的数据集管理综合解决方案。它使您能够跟踪、版本化和管理用于训练、验证和评估的数据集，提供从原始数据到模型预测的完整血缘关系。

为何数据集跟踪如此重要​

数据集跟踪对于可复现的机器学习至关重要，并提供以下几个关键优势：

  • 数据血缘：跟踪从原始数据源到模型输入的完整过程。
  • 可复现性：确保实验可以使用完全相同的数据集进行复现。
  • 版本控制：管理数据集随其演变的不同版本。
  • 协作：在团队间共享数据集及其元数据。
  • 评估集成：与 MLflow 的评估功能无缝集成。
  • 生产监控：跟踪在生产推理和评估中使用的数据集。

核心组件​

MLflow 的数据集跟踪围绕两个主要抽象展开：

Dataset (数据集)​

Dataset 抽象是一个元数据跟踪对象，它持有关于已记录数据集的全面信息。存储在 Dataset 对象中的信息包括：

核心属性

  • Name (名称)：数据集的描述性标识符（如果未指定，则默认为“dataset”）。
  • Digest (摘要)：用于数据集识别的唯一哈希/指纹（自动计算）。
  • Source (来源)：包含指向原始数据位置的血缘信息的 DatasetSource。
  • Schema (模式)：可选的数据集模式（具体实现相关，例如 MLflow Schema）。
  • Profile (概览)：可选的摘要统计信息（具体实现相关，例如行数、列统计信息）。

支持的数据集类型

特殊数据集类型

  • EvaluationDataset - 一种内部数据集类型，专门与 mlflow.evaluate() 一起用于模型评估工作流。

DatasetSource (数据集来源)​

DatasetSource 组件提供了与数据原始来源的链接血缘关系，无论是文件 URL、S3 存储桶、数据库表还是任何其他数据源。这确保您始终可以追溯到数据的来源。

DatasetSource 可以使用 mlflow.data.get_source() API 来检索，该 API 接受 Dataset、DatasetEntity 或 DatasetInput 的实例。

快速入门：基本数据集跟踪​

以下是如何开始基本数据集跟踪的方法：

import mlflow.data
import pandas as pd

# Load your data
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 a Dataset object
dataset = mlflow.data.from_pandas(
    raw_data, source=dataset_source_url, name="wine-quality-white", targets="quality"
)

# Log the dataset to an MLflow run
with mlflow.start_run():
    mlflow.log_input(dataset, context="training")

    # Your training code here
    # model = train_model(raw_data)
    # mlflow.sklearn.log_model(model, "model")

数据集信息和元数据​

当您创建数据集时，MLflow 会自动捕获丰富的元数据：

# Access dataset metadata
print(f"Dataset name: {dataset.name}")  # Defaults to "dataset" if not specified
print(
    f"Dataset digest: {dataset.digest}"
)  # Unique hash identifier (computed automatically)
print(f"Dataset source: {dataset.source}")  # DatasetSource object
print(
    f"Dataset profile: {dataset.profile}"
)  # Optional: implementation-specific statistics
print(f"Dataset schema: {dataset.schema}")  # Optional: implementation-specific schema

示例输出：

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": "long", "name": "quality"}
]}
Dataset source: <DatasetSource object>
数据集属性

profile 和 schema 属性是具体实现相关的，可能会根据数据集类型（PandasDataset、SparkDataset 等）而有所不同。某些数据集类型可能会为这些属性返回 None。

数据集来源和血缘​

MLflow 支持来自多种来源的数据集：

# From local file
local_dataset = mlflow.data.from_pandas(
    df, source="/path/to/local/file.csv", name="local-data"
)

# From cloud storage
s3_dataset = mlflow.data.from_pandas(
    df, source="s3://bucket/data.parquet", name="s3-data"
)

# From database
db_dataset = mlflow.data.from_pandas(
    df, source="postgresql://user:pass@host/db", name="db-data"
)

# From URL
url_dataset = mlflow.data.from_pandas(
    df, source="https://example.com/data.csv", name="web-data"
)

MLflow UI 中的数据集跟踪​

当您将数据集记录到 MLflow 运行中时，它们会与全面的元数据一起出现在 MLflow UI 中。您可以直接在界面中查看数据集信息、模式和血缘。

Dataset in MLflow UI

UI 显示：

  • 数据集名称和摘要
  • 带列类型的模式信息
  • 概览统计信息（行数等）
  • 来源血缘信息
  • 数据集使用的上下文

与 MLflow Evaluate 集成​

MLflow 数据集最强大的功能之一是它们与 MLflow 评估功能的无缝集成。在使用 mlflow.evaluate() 时，MLflow 会在内部自动将各种数据类型转换为 EvaluationDataset 对象。

EvaluationDataset

MLflow 在使用 mlflow.evaluate() 时会使用一个内部的 EvaluationDataset 类。此数据集类型是根据您的输入数据自动创建的，并为评估工作流提供优化的哈希和元数据跟踪。

直接使用数据集与 MLflow evaluate

import mlflow
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split

# Prepare data and train model
data = pd.read_csv("classification_data.csv")
X = data.drop("target", axis=1)
y = data["target"]
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42
)

model = RandomForestClassifier()
model.fit(X_train, y_train)

# Create evaluation dataset
eval_data = X_test.copy()
eval_data["target"] = y_test

eval_dataset = mlflow.data.from_pandas(
    eval_data, targets="target", name="evaluation-set"
)

with mlflow.start_run():
    # Log model
    mlflow.sklearn.log_model(model, name="model", input_example=X_test)

    # Evaluate using the dataset
    result = mlflow.evaluate(
        model="runs:/{}/model".format(mlflow.active_run().info.run_id),
        data=eval_dataset,
        model_type="classifier",
    )

    print(f"Accuracy: {result.metrics['accuracy_score']:.3f}")

MLflow Evaluate 集成示例​

以下是一个完整的示例，展示了数据集如何与 MLflow 的评估功能集成：

Dataset Evaluation in MLflow UI

评估运行展示了数据集、模型、指标和评估产物（如混淆矩阵）如何全部被记录在一起，提供了评估过程的完整视图。

高级数据集管理​

跟踪数据集随其演变的版本

def create_versioned_dataset(data, version, base_name="customer-data"):
    """Create a versioned dataset with metadata."""

    dataset = mlflow.data.from_pandas(
        data,
        source=f"data_pipeline_v{version}",
        name=f"{base_name}-v{version}",
        targets="target",
    )

    with mlflow.start_run(run_name=f"Dataset_Version_{version}"):
        mlflow.log_input(dataset, context="versioning")

        # Log version metadata
        mlflow.log_params(
            {
                "dataset_version": version,
                "data_size": len(data),
                "features_count": len(data.columns) - 1,
                "target_distribution": data["target"].value_counts().to_dict(),
            }
        )

        # Log data quality metrics
        mlflow.log_metrics(
            {
                "missing_values_pct": (data.isnull().sum().sum() / data.size) * 100,
                "duplicate_rows": data.duplicated().sum(),
                "target_balance": data["target"].std(),
            }
        )

    return dataset


# Create multiple versions
v1_dataset = create_versioned_dataset(data_v1, "1.0")
v2_dataset = create_versioned_dataset(data_v2, "2.0")
v3_dataset = create_versioned_dataset(data_v3, "3.0")

生产用例​

监控生产批量预测中使用的数据集

def monitor_batch_predictions(batch_data, model_version, date):
    """Monitor production batch prediction datasets."""

    # Create dataset for batch predictions
    batch_dataset = mlflow.data.from_pandas(
        batch_data,
        source=f"production_batch_{date}",
        name=f"batch_predictions_{date}",
        targets="true_label" if "true_label" in batch_data.columns else None,
        predictions="prediction" if "prediction" in batch_data.columns else None,
    )

    with mlflow.start_run(run_name=f"Batch_Monitor_{date}"):
        mlflow.log_input(batch_dataset, context="production_batch")

        # Log production metadata
        mlflow.log_params(
            {
                "batch_date": date,
                "model_version": model_version,
                "batch_size": len(batch_data),
                "has_ground_truth": "true_label" in batch_data.columns,
            }
        )

        # Monitor prediction distribution
        if "prediction" in batch_data.columns:
            pred_metrics = {
                "prediction_mean": batch_data["prediction"].mean(),
                "prediction_std": batch_data["prediction"].std(),
                "unique_predictions": batch_data["prediction"].nunique(),
            }
            mlflow.log_metrics(pred_metrics)

        # Evaluate if ground truth is available
        if all(col in batch_data.columns for col in ["prediction", "true_label"]):
            result = mlflow.evaluate(data=batch_dataset, model_type="classifier")
            print(f"Batch accuracy: {result.metrics.get('accuracy_score', 'N/A')}")

    return batch_dataset


# Usage
batch_dataset = monitor_batch_predictions(daily_batch_data, "v2.1", "2024-01-15")

最佳实践​

在使用 MLflow 数据集时，请遵循以下最佳实践：

数据质量：在记录数据集之前始终验证数据质量。检查缺失值、重复项和数据类型。

命名约定：为数据集使用一致的、描述性的名称，其中包含版本信息和上下文。

来源文档：始终指定有意义的源 URL 或标识符，以便您可以追溯到原始数据。

上下文规范：在记录数据集时使用清晰的上下文标签（例如，“training”、“validation”、“evaluation”、“production”）。

元数据记录：包含有关数据收集、预处理步骤和数据特征的相关元数据。

版本控制：明确跟踪数据集版本，尤其是在数据预处理或收集方法发生变化时。

摘要计算：数据集摘要的计算方式因数据集类型而异：

  • 标准数据集：基于数据内容和结构。
  • MetaDataset：仅基于元数据（名称、来源、模式）- 不进行实际数据哈希。
  • EvaluationDataset：针对大型数据集使用样本行进行优化哈希。

来源灵活性：DatasetSource 支持多种来源类型，包括 HTTP URL、文件路径、数据库连接和云存储位置。

评估集成：通过明确指定目标列和预测列来设计数据集，以便于评估。

主要优势​

MLflow 数据集跟踪为 ML 团队提供了几个关键优势：

可复现性：确保即使数据源演变，实验也可以使用完全相同的数据集进行复现。

血缘跟踪：维护从源到模型预测的完整数据血缘，从而实现更好的调试和合规性。

协作：通过一致的界面在团队成员之间共享数据集及其元数据。

评估集成：与 MLflow 的评估功能无缝集成，以进行全面的模型评估。

生产监控：跟踪生产系统中用于性能监控和数据漂移检测的数据集。

质量保证：自动捕获数据质量指标并随时间监控变化。

无论您是跟踪训练数据集、管理评估数据，还是监控生产批量预测，MLflow 的数据集跟踪功能都为可靠、可复现的机器学习工作流提供了基础。