ZenML项目教程：使用自定义数据集类管理机器学习数据

ZenML项目教程：使用自定义数据集类管理机器学习数据

zenml ZenML 🙏: Build portable, production-ready MLOps pipelines. https://zenml.io. 项目地址: https://gitcode.com/gh_mirrors/ze/zenml

在机器学习项目中，随着项目复杂度的增加，我们经常需要处理来自不同来源的数据并管理复杂的数据流。本文将详细介绍如何在ZenML框架中使用自定义数据集类(Dataset)和材质化器(Materializer)来高效处理这些挑战。

为什么需要自定义数据集类？

在传统机器学习项目中，数据加载和处理逻辑往往分散在各个脚本中，导致代码难以维护和扩展。ZenML通过自定义数据集类提供了以下优势：

1. 统一接口：为不同数据源提供一致的读取和写入接口
2. 封装复杂性：将特定数据源的实现细节隐藏在类内部
3. 增强可维护性：集中管理数据相关逻辑，便于修改和扩展
4. 支持多种数据源：轻松集成CSV、数据库、云存储等多种数据源

自定义数据集类实现详解

基础抽象类设计

首先，我们定义一个抽象基类，规定所有数据集类必须实现的基本方法：

from abc import ABC, abstractmethod
import pandas as pd
from typing import Optional

class Dataset(ABC):
    @abstractmethod
    def read_data(self) -> pd.DataFrame:
        """读取数据并返回DataFrame"""
        pass

这个简单的接口确保了所有具体数据集类都提供统一的数据访问方式。

CSV数据集实现

对于CSV文件，我们可以实现如下具体类：

class CSVDataset(Dataset):
    def __init__(self, data_path: str, df: Optional[pd.DataFrame] = None):
        self.data_path = data_path  # CSV文件路径
        self.df = df  # 可选：直接传入DataFrame避免重复读取

    def read_data(self) -> pd.DataFrame:
        if self.df is None:
            self.df = pd.read_csv(self.data_path)
        return self.df

这个实现支持两种数据加载方式：

1. 从指定路径读取CSV文件
2. 直接使用传入的DataFrame（适用于已经在内存中的数据）

BigQuery数据集实现

对于Google BigQuery数据源，实现会稍微复杂一些：

from google.cloud import bigquery

class BigQueryDataset(Dataset):
    def __init__(
        self,
        table_id: str,
        df: Optional[pd.DataFrame] = None,
        project: Optional[str] = None,
    ):
        self.table_id = table_id  # BigQuery表ID
        self.project = project  # GCP项目ID
        self.df = df  # 可选：缓存DataFrame
        self.client = bigquery.Client(project=self.project)

    def read_data(self) -> pd.DataFrame:
        query = f"SELECT * FROM `{self.table_id}`"
        self.df = self.client.query(query).to_dataframe()
        return self.df

    def write_data(self) -> None:
        """将DataFrame写回BigQuery"""
        job_config = bigquery.LoadJobConfig(write_disposition="WRITE_TRUNCATE")
        job = self.client.load_table_from_dataframe(
            self.df, self.table_id, job_config=job_config
        )
        job.result()

这个实现不仅支持数据读取，还提供了数据写回BigQuery的功能。

材质化器(Materializer)实现原理

材质化器是ZenML中负责对象序列化和反序列化的组件。对于自定义数据集类，我们需要实现专门的材质化器。

CSV数据集材质化器

from zenml.materializers import BaseMaterializer
import tempfile
import os

class CSVDatasetMaterializer(BaseMaterializer):
    ASSOCIATED_TYPES = (CSVDataset,)  # 关联的数据集类型
    ASSOCIATED_ARTIFACT_TYPE = ArtifactType.DATA  # 关联的工件类型

    def load(self, data_type: Type[CSVDataset]) -> CSVDataset:
        # 创建临时文件保存CSV数据
        with tempfile.NamedTemporaryFile(delete=False, suffix='.csv') as temp_file:
            # 从工件存储复制到临时位置
            with fileio.open(os.path.join(self.uri, "data.csv"), "rb") as source_file:
                temp_file.write(source_file.read())
            
            temp_path = temp_file.name

        # 创建并返回CSVDataset实例
        dataset = CSVDataset(temp_path)
        dataset.read_data()  # 预加载数据
        return dataset

    def save(self, dataset: CSVDataset) -> None:
        # 确保有数据可保存
        df = dataset.read_data()

        # 保存到临时CSV文件
        with tempfile.NamedTemporaryFile(delete=False, suffix='.csv') as temp_file:
            df.to_csv(temp_file.name, index=False)
            temp_path = temp_file.name

        # 复制到工件存储
        with open(temp_path, "rb") as source_file:
            with fileio.open(os.path.join(self.uri, "data.csv"), "wb") as target_file:
                target_file.write(source_file.read())

        # 清理临时文件
        os.remove(temp_path)

BigQuery数据集材质化器

import json

class BigQueryDatasetMaterializer(BaseMaterializer):
    ASSOCIATED_TYPES = (BigQueryDataset,)
    ASSOCIATED_ARTIFACT_TYPE = ArtifactType.DATA

    def load(self, data_type: Type[BigQueryDataset]) -> BigQueryDataset:
        # 加载元数据
        with fileio.open(os.path.join(self.uri, "metadata.json"), "r") as f:
            metadata = json.load(f)
        
        # 重建BigQueryDataset实例
        dataset = BigQueryDataset(
            table_id=metadata["table_id"],
            project=metadata["project"],
        )
        dataset.read_data()  # 预加载数据
        return dataset

    def save(self, bq_dataset: BigQueryDataset) -> None:
        # 保存元数据
        metadata = {
            "table_id": bq_dataset.table_id,
            "project": bq_dataset.project,
        }
        with fileio.open(os.path.join(self.uri, "metadata.json"), "w") as f:
            json.dump(metadata, f)
        
        # 如果有数据，写回BigQuery
        if bq_dataset.df is not None:
            bq_dataset.write_data()

构建灵活的数据处理流水线

有了自定义数据集类和材质化器，我们可以构建灵活处理多种数据源的流水线：

from zenml import step, pipeline
from typing_extensions import Annotated

@step(output_materializer=CSVDatasetMaterializer)
def extract_data_local(data_path: str = "data/raw_data.csv") -> CSVDataset:
    """本地CSV数据提取步骤"""
    return CSVDataset(data_path)

@step(output_materializer=BigQueryDatasetMaterializer)
def extract_data_remote(table_id: str) -> BigQueryDataset:
    """远程BigQuery数据提取步骤"""
    return BigQueryDataset(table_id)

@step
def transform(dataset: Dataset) -> pd.DataFrame:
    """通用数据转换步骤"""
    df = dataset.read_data()
    # 这里添加实际的数据转换逻辑
    transformed_df = df.copy()  # 示例中仅复制，实际应添加转换
    return transformed_df

@pipeline
def etl_pipeline(mode: str = "develop"):
    """ETL流水线，根据模式选择数据源"""
    if mode == "develop":
        raw_data = extract_data_local()
    else:
        raw_data = extract_data_remote(table_id="project.dataset.raw_table")

    transformed_data = transform(raw_data)

这种设计允许我们在开发和部署阶段使用不同的数据源，而无需修改核心处理逻辑。

最佳实践总结

1. 分层抽象：使用基类定义统一接口，具体类处理特定数据源细节
2. 关注点分离：数据加载、转换和存储逻辑应分开处理
3. 灵活配置：通过参数控制数据源选择，而非硬编码
4. 资源管理：妥善处理临时文件和外部连接
5. 缓存优化：避免重复读取大数据集
6. 错误处理：为不同数据源添加适当的错误恢复机制

通过遵循这些模式，你可以构建出既灵活又可靠的机器学习数据流水线，轻松应对各种复杂的数据处理场景。

zenml ZenML 🙏: Build portable, production-ready MLOps pipelines. https://zenml.io. 项目地址: https://gitcode.com/gh_mirrors/ze/zenml