ZenML项目实战：如何创建自定义数据可视化组件-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00718/article/details/148524483

ZenML项目实战：如何创建自定义数据可视化组件

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

在机器学习工作流中，数据可视化是理解模型行为和数据分析结果的关键环节。ZenML作为一个开源的机器学习运维(MLOps)框架，提供了灵活的可视化机制，允许开发者根据项目需求创建自定义的可视化组件。本文将深入探讨ZenML中的三种可视化实现方式，并通过实际案例展示如何为特定数据类型构建专属可视化方案。

一、可视化基础概念

在ZenML架构中，可视化功能与数据工件(Artifact)紧密关联。每个步骤(Step)产生的输出都可以通过多种方式呈现可视化效果，这些可视化结果会直接展示在ZenML的仪表盘(Dashboard)中，为团队提供直观的数据洞察。

二、三种可视化实现方式

2.1 使用特殊返回类型实现可视化

适用场景：当步骤中已经生成了HTML、Markdown或CSV格式的字符串数据时，这是最快捷的可视化实现方式。

ZenML提供了三种特殊类型包装器：

HTMLString：用于HTML格式字符串
MarkdownString：用于Markdown格式字符串
CSVString：用于CSV格式字符串

实现示例：

from zenml.types import HTMLString

@step
def generate_html_report() -> HTMLString:
    report = """
    <html>
        <body>
            <h1>模型训练报告</h1>
            <p>准确率: 95%</p>
            <p>召回率: 93%</p>
        </body>
    </html>
    """
    return HTMLString(report)

技术要点：

返回类型必须明确声明为对应的特殊类型
字符串内容必须符合对应格式规范
系统会自动识别类型并渲染相应可视化组件

2.2 通过Materializer实现类型可视化

适用场景：需要为特定数据类型(如Pandas DataFrame、NumPy数组等)自动生成可视化时。

Materializer是ZenML中负责数据序列化和反序列化的组件，通过重写其save_visualizations方法可以实现自动化可视化。

实现步骤：

继承BaseMaterializer创建自定义Materializer
在ASSOCIATED_TYPES中关联目标数据类型
实现save_visualizations方法生成可视化文件

优势：

一次开发，多处使用
自动应用于所有匹配类型的工件
支持复杂可视化逻辑

2.3 自定义返回类型结合Materializer

适用场景：需要实现高度定制化可视化方案时。

这是最灵活的方式，结合了前两种方法的优点，适合构建专业级的可视化组件。

实现流程：

创建自定义数据类（继承BaseModel）
开发专属Materializer并实现可视化逻辑
在步骤中返回自定义类实例

三、实战案例：Facets数据对比可视化

以Google的Facets工具集成示例，展示完整实现过程：

3.1 定义数据模型

from pydantic import BaseModel
import pandas as pd
from typing import Dict, List, Union

class FacetsComparison(BaseModel):
    """存储待对比的数据集信息"""
    datasets: List[Dict[str, Union[str, pd.DataFrame]]]

3.2 实现Materializer

from zenml.materializers.base_materializer import BaseMaterializer
from zenml.artifact_stores.models import ArtifactType
import os
from zenml.io import fileio

VISUALIZATION_FILENAME = "facets_overview.html"

class FacetsMaterializer(BaseMaterializer):
    ASSOCIATED_TYPES = (FacetsComparison,)
    ASSOCIATED_ARTIFACT_TYPE = ArtifactType.DATA_ANALYSIS

    def save_visualizations(self, data: FacetsComparison) -> Dict[str, str]:
        # 生成Facets可视化HTML
        html = generate_facets_html(data.datasets)
        
        # 保存可视化文件
        viz_path = os.path.join(self.uri, VISUALIZATION_FILENAME)
        with fileio.open(viz_path, "w") as f:
            f.write(html)
        
        return {viz_path: "text/html"}

3.3 创建业务步骤

@step
def compare_datasets(
    training_data: pd.DataFrame,
    production_data: pd.DataFrame,
) -> FacetsComparison:
    """对比训练数据和生产环境数据的分布"""
    return FacetsComparison(
        datasets=[
            {"name": "训练数据", "table": training_data},
            {"name": "生产数据", "table": production_data}
        ]
    )