ZenML项目实战：自定义数据类型处理与Materializer开发指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00162/article/details/148524250

ZenML项目实战：自定义数据类型处理与Materializer开发指南

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

什么是Materializer？

在ZenML项目中，Materializer（物化器）是一个核心组件，它负责管理数据在流水线步骤间的传递方式。简单来说，Materializer定义了如何将特定类型的数据序列化存储到Artifact Store（工件存储），以及如何从存储中反序列化读取数据。

为什么需要Materializer？

在机器学习流水线中，每个步骤都是独立的执行单元，它们通过输入和输出相互连接。Materializer的作用在于：

数据持久化：将步骤输出的数据保存到持久化存储
数据加载：在后续步骤中重新加载这些数据
类型转换：确保数据在不同步骤间保持正确的类型

内置Materializer一览

ZenML已经为常见数据类型提供了开箱即用的Materializer支持：

| 数据类型 | 存储格式 | 说明 | |---------|---------|------| | 基本类型(bool, float等) | JSON | 处理Python基础数据类型 | | NumPy数组 | .npy | 专为NumPy数组优化 | | Pandas DataFrame | CSV/Parquet | 支持大数据集的高效存储 | | Pydantic模型 | JSON | 完美支持Pydantic数据模型 |

集成Materializer扩展

除了内置支持，ZenML还通过集成提供了更多专业领域的Materializer：

深度学习框架：PyTorch、TensorFlow、HuggingFace模型
数据分析工具：Spark、Polars、Great Expectations
模型服务：BentoML、PyCaret

这些集成Materializer需要通过安装对应扩展包来激活。

开发自定义Materializer

当内置Materializer不能满足需求时，我们可以创建自定义实现。以下是关键步骤：

1. 基础类继承

所有Materializer必须继承自BaseMaterializer基类：

from zenml.materializers.base_materializer import BaseMaterializer

class CustomMaterializer(BaseMaterializer):
    ASSOCIATED_TYPES = (CustomType,)  # 指定处理的数据类型
    ASSOCIATED_ARTIFACT_TYPE = ArtifactType.DATA  # 指定工件类型

2. 实现核心方法

必须实现load和save两个核心方法：

def save(self, data: CustomType) -> None:
    """实现数据序列化存储逻辑"""
    # 示例：将自定义对象保存为JSON
    with self.artifact_store.open(self.uri, "w") as f:
        json.dump(data.to_dict(), f)

def load(self, data_type: Type[CustomType]) -> CustomType:
    """实现数据反序列化逻辑"""
    # 示例：从JSON重建对象
    with self.artifact_store.open(self.uri, "r") as f:
        return data_type.from_dict(json.load(f))

3. 注册Materializer

有两种注册方式：

方式一：步骤级注册

@step(output_materializers=CustomMaterializer)
def processing_step() -> CustomType:
    ...

方式二：全局注册

from zenml.materializers.materializer_registry import materializer_registry

materializer_registry.register_and_overwrite_type(
    key=CustomType,
    type_=CustomMaterializer
)

最佳实践建议

避免使用Cloudpickle：虽然方便，但存在版本兼容性和安全隐患
明确指定Materializer：不要依赖自动发现，提高代码可维护性
考虑存储效率：大数据集优先使用二进制格式而非文本
类型注解要清晰：使用Python类型提示提高代码可读性

实战示例：自定义图像处理器

假设我们需要处理特殊图像格式，可以这样实现：

from PIL import Image
from zenml.materializers import BaseMaterializer

class SpecialImageMaterializer(BaseMaterializer):
    ASSOCIATED_TYPES = (SpecialImage,)
    ASSOCIATED_ARTIFACT_TYPE = ArtifactType.DATA

    def save(self, image: SpecialImage) -> None:
        """将特殊图像转换为PNG存储"""
        temp_path = f"{self.uri}.png"
        image.to_pil().save(temp_path)
        self.artifact_store.upload(temp_path, self.uri)

    def load(self, data_type: Type[SpecialImage]) -> SpecialImage:
        """从PNG重建特殊图像"""
        temp_path = f"{self.uri}.png"
        self.artifact_store.download(self.uri, temp_path)
        return data_type.from_pil(Image.open(temp_path))