ZenML项目实战：从零开始构建机器学习流水线

ZenML项目实战：从零开始构建机器学习流水线

zenml 项目地址: https://gitcode.com/gh_mirrors/zen/zenml

前言

在机器学习项目开发过程中，如何将数据处理、模型训练等环节组织成可复用的工作流是一个常见挑战。本文将介绍如何使用ZenML框架构建端到端的机器学习流水线，通过模块化设计提高项目的可维护性和可扩展性。

什么是ZenML

ZenML是一个开源的机器学习操作(MLOps)框架，它通过标准化的方式组织机器学习工作流。核心概念包括：

• 步骤(Step)：代表工作流中的一个独立单元，如数据加载、特征工程等
• 流水线(Pipeline)：将多个步骤连接起来形成完整工作流
• 工件(Artifact)：步骤之间传递的数据对象

环境准备

在开始前，请确保已安装ZenML：

pip install "zenml[server]"
zenml up  # 启动本地仪表盘

第一个简单流水线

让我们从一个基础示例开始，了解ZenML的基本用法：

from zenml import pipeline, step

@step
def load_data() -> dict:
    """模拟加载训练数据和标签"""
    training_data = [[1, 2], [3, 4], [5, 6]]
    labels = [0, 1, 0]
    return {'features': training_data, 'labels': labels}

@step
def train_model(data: dict) -> None:
    """模拟模型训练过程"""
    total_features = sum(map(sum, data['features']))
    total_labels = sum(data['labels'])
    print(f"使用{len(data['features'])}个数据点训练模型。特征总和{total_features}，标签总和{total_labels}")

@pipeline
def simple_ml_pipeline():
    """定义连接各步骤的流水线"""
    dataset = load_data()
    train_model(dataset)

if __name__ == "__main__":
    run = simple_ml_pipeline()

关键点说明：

• @step装饰器将普通函数转换为流水线步骤
• @pipeline装饰器定义整个工作流
• 步骤间通过返回值自动传递数据

执行后可以在仪表盘中查看运行结果，直观了解各步骤执行情况和数据流向。

完整机器学习工作流示例

下面我们构建一个更实际的Iris分类器流水线：

1. 安装额外依赖

pip install matplotlib
zenml integration install sklearn -y

2. 数据加载步骤

from typing import Tuple, Annotated
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
import logging

@step
def training_data_loader() -> Tuple[
    Annotated[pd.DataFrame, "X_train"],
    Annotated[pd.DataFrame, "X_test"],
    Annotated[pd.Series, "y_train"],
    Annotated[pd.Series, "y_test"],
]:
    """加载Iris数据集并分割"""
    logging.info("加载数据中...")
    iris = load_iris(as_frame=True)
    X_train, X_test, y_train, y_test = train_test_split(
        iris.data, iris.target, test_size=0.2, random_state=42
    )
    return X_train, X_test, y_train, y_test

3. 模型训练步骤

from sklearn.base import ClassifierMixin
from sklearn.svm import SVC

@step
def svc_trainer(
    X_train: pd.DataFrame,
    y_train: pd.Series,
    gamma: float = 0.001,
) -> Tuple[
    Annotated[ClassifierMixin, "trained_model"],
    Annotated[float, "training_acc"],
]:
    """训练SVC分类器"""
    model = SVC(gamma=gamma)
    model.fit(X_train.to_numpy(), y_train.to_numpy())
    train_acc = model.score(X_train.to_numpy(), y_train.to_numpy())
    print(f"训练准确率: {train_acc}")
    return model, train_acc

4. 组合成完整流水线

@pipeline
def training_pipeline(gamma: float = 0.002):
    """完整训练流水线"""
    X_train, X_test, y_train, y_test = training_data_loader()
    svc_trainer(gamma=gamma, X_train=X_train, y_train=y_train)

if __name__ == "__main__":
    training_pipeline(gamma=0.0015)

使用YAML配置

为方便参数调整，可以将配置移至YAML文件：

# config.yaml
parameters:
    gamma: 0.01

然后在代码中引用：

training_pipeline = training_pipeline.with_options(
    config_path='config.yaml'
)
training_pipeline()

最佳实践建议

1. 模块化设计：将不同功能拆分为独立步骤
2. 类型注解：明确输入输出数据类型
3. 日志记录：使用logging模块而非print
4. 配置分离：将超参数等移至配置文件
5. 版本控制：检查配置文件到代码仓库

总结

通过ZenML，我们可以将机器学习工作流标准化为可复用的组件，提高开发效率和项目可维护性。本文展示了从简单到完整的流水线构建过程，读者可以在此基础上扩展更复杂的场景，如模型评估、部署等环节。

ZenML的自动跟踪和版本控制功能也为团队协作和模型迭代提供了良好基础，是构建生产级机器学习应用的有力工具。

zenml 项目地址: https://gitcode.com/gh_mirrors/zen/zenml