ZenML + Optuna 超参数调优实战：基于动态管线与 Ask API 的并行试验编排

ZenML Optuna 超参数调优实战基于动态管线与 Ask API 的并行试验编排【免费下载链接】zenmlZenML : One AI Platform from Pipelines to Agents. https://zenml.io.项目地址: https://gitcode.com/GitHub_Trending/ze/zenml本文以 examples/optuna_hyperparameter_tuning 为骨架系统讲解如何将 Optuna 的 ask API 与 ZenML 动态管线pipeline(dynamicTrue)step.map()结合在 FashionMNIST 上对 PyTorch Lightning CNN 进行智能超参数搜索。读完本文你将掌握简单并行扫描与自适应多轮优化两种模式的设计思路、完整运行方式以及如何借助 ZenML 构件artifact流转试验历史让超参数搜索在本地与分布式编排器上均可移植、可复现、可查询。核心思路Optuna 决定试什么ZenML 决定在哪跑超参数搜索本质上是采样一批配置 → 训练 → 评估 → 反馈给采样器的闭环。传统做法通常依赖一个共享的 Optuna 存储后端如 SQLite/MySQL在多机、多容器场景下引入额外运维负担。本示例的关键设计是Optuna study 以内存方式在每个suggest_trials步骤内创建历史试验结果通过 ZenML 构件artifact而非共享数据库在轮次之间传递从而彻底解耦智能采样与任务编排两个关注点Optuna 负责智能采样通过 ask API 生成下一批值得尝试的超参数组合TPE 采样器会从历史结果中学习ZenML 负责执行编排通过动态管线把每个试验扇出fan out为独立步骤并行执行统一管理并行度、缓存、资源池与元数据记录。该架构的 ASCII 流程如下摘自 READMEOptuna Study (In-Memory) | - Ask API: Generate trial configs | | | - ZenML Pipeline: Orchestrate parallel training | | | | | -- train_trial_0 -- | | -- train_trial_1 -- Parallel execution | | -- train_trial_2 | (managed by orchestrator) | | -- train_trial_N -- | | | | ---------------- Collect results | - Artifact History: Results flow through ZenML artifacts (no shared DB needed)关键结论试验历史流经 ZenML 构件因此这套方案在本地与分布式编排器如 Kubernetes上都能工作无需共享数据库。快速开始环境准备示例依赖清单见 requirements.txt核心依赖包括zenml[server]0.93.1、optuna3.0.0、filelock3.0.0以及torch2.0.0、lightning2.0.0、torchvision0.15.0。初始化环境uv venv --seed source .venv/bin/activate pip install -r requirements.txt zenml init zenml login如果你使用基于 AWS Kubernetes 的 stack还需要安装对应 ZenML 集成zenml integration install aws s3 kubernetes --yes --uv运行简单扫描Simple Sweep所有试验并行执行python run.py # ^^^ 默认使用 config/simple.yaml —— 2 个并行试验每个训练 3 个 epoch # 也可以显式指定模式或自定义配置 python run.py simple python run.py simple --config config/my_experiment.yaml # 自定义配置运行自适应扫描Adaptive Sweep多轮执行每一轮都从上一轮结果中学习python run.py adaptive python run.py adaptive --config config/my_experiment.yaml默认使用config/adaptive.yaml—— 2 轮、每轮 2 个试验共 4 个试验、每个试验 3 个 epoch。Optuna 的 TPE 采样器利用早期结果在后几轮给出更优的超参数建议。命令行入口解析run.py 是唯一入口使用 argparse 提供三个参数参数取值/默认值说明modesimple默认/adaptive选择简单并行扫描或自适应多轮扫描--config默认config/simple.yaml或config/adaptive.yaml指定 YAML 配置文件的路径--no-cache默认开启缓存关闭管线缓存强制重新执行步骤例如python run.py adaptive --no-cache表示以自适应模式、关闭缓存运行。入口通过pipeline.with_options(config_path..., enable_cache...)将配置注入对应管线见 run.py。配置文件详解每种模式有独立的 YAML 配置包含管线参数与 Docker 设置。config/simple.yaml对应sweep_pipeline参数n_trials、max_iter等config/adaptive.yaml对应adaptive_sweep_pipeline参数n_rounds、trials_per_round、max_iter等。以 config/simple.yaml 为例enable_step_logs: true enable_artifact_metadata: true enable_artifact_visualization: true parameters: study_name: fashion_mnist_sweep n_trials: 2 max_iter: 3 settings: docker: requirements: requirements.txt python_package_installer: uv build_config: build_options: platform: linux/amd64parameters直接映射到管线函数的同名入参study_name、n_trials、max_iter用于控制试验数量、训练 epoch 数等enable_*三个开关控制步骤日志、构件元数据与构件可视化settings.docker定义远程执行时的镜像构建方式依赖来自requirements.txt包管理器使用uv目标平台为linux/amd64。config/adaptive.yaml 结构相同仅parameters换成了n_rounds: 2、trials_per_round: 2。编辑这些文件即可调整试验规模、训练轮数、缓存行为与远程执行时的 Docker 设置。你会看到什么Dashboard 中的两种 DAGSimple Sweep DAGsuggest_trials (generates 5 configs) | -- train_trial_0 -- -- train_trial_1 | -- train_trial_2 -- All run in parallel -- train_trial_3 | (managed by orchestrator) -- train_trial_4 -- | report_results (aggregates metrics, finds best) | retrain_best_model (saves production model)关键设计试验模型被丢弃只保存指标。找到最优超参数后管线用该参数重新训练只保存这一个生产就绪的模型避免存储浪费。Adaptive Sweep DAGRound 1: suggest_trials_round_0 -- train_trial_0, 1 (parallel) -- report_results_round_0 Round 2: suggest_trials_round_1 (learns from round 1) -- train_trial_2, 3 (parallel) -- report_results_round_1 Round 3: suggest_trials_round_2 (learns from rounds 1 2) -- train_trial_4, 5 (parallel) -- report_results_round_2 | retrain_best_model (saves production model)元数据与构件每个试验构件都附带丰富元数据真实结构来自 steps/train.py 的log_metadata调用{ trial_number: 3, val_loss: 0.4321, val_accuracy: 84.52, learning_rate: 0.003421, batch_size: 64, hidden_dim: 32, n_epochs: 10, }这使得试验结果可查询Dashboard按val_loss、learning_rate等字段过滤、排序MCP 服务器用自然语言查询如show me all trials with val_loss 0.4ZenML Clientclient.list_artifact_versions()配合元数据过滤器编程式查询。深入原理五个步骤的源码拆解1. 建议试验steps/suggest.py该步骤使用 Optunaask API生成试验配置。关键点在于内存 study 每次调用时创建并通过 ZenML 构件把上一轮结果重放进去让 TPE 采样器基于历史学习study optuna.create_study(study_namestudy_name, directionminimize) # 重放历史结果使 TPE 采样器从中学习 if previous_summary and previous_summary.get(all_trials): distributions _build_distributions(SEARCH_SPACE) for trial_info in previous_summary[all_trials]: study.add_trial( optuna.trial.create_trial( params{...}, distributionsdistributions, values[trial_info[val_loss]], stateoptuna.trial.TrialState.COMPLETE, ) ) # 用 ask() 生成新试验配置 for _ in range(n_trials): trial study.ask() config { trial_number: trial.number, learning_rate: trial.suggest_float( learning_rate, 1e-4, 1e-2, logTrue ), batch_size: trial.suggest_categorical(batch_size, [64, 128]), hidden_dim: trial.suggest_categorical(hidden_dim, [8, 16, 32]), } trials.append(config)返回一个配置列表每个试验一条。值得注意的实现细节_build_distributions会把SEARCH_SPACE字典翻译成 Optuna 的FloatDistribution对数均匀与CategoricalDistribution且重放历史时必须复用同一组分布对象否则会出现动态值空间冲突见下文排障。该步骤通过settings声明cpu_count1、memory1GB的资源需求。2. 并行训练steps/train.py每个试验用 PyTorch Lightning 独立训练一个 CNN互不干扰从而支持并行step def train_trial(trial_config: dict, max_iter: int 10) - dict: model FashionMNISTClassifier( learning_ratetrial_config[learning_rate], hidden_dimtrial_config[hidden_dim], ) trainer Trainer( max_epochsmax_iter, acceleratorauto, devices1, callbacks[EarlyStopping(monitorval_loss, patience10, modemin)], enable_progress_barFalse, loggerFalse, deterministicTrue, ) trainer.fit(model, train_loader, val_loader) log_metadata({ trial_number: trial_config[trial_number], val_loss: val_loss, val_accuracy: val_accuracy, ... }, artifact_nametrial_result, infer_artifactTrue) return {trial_number: ..., val_loss: val_loss, ...}源码层面的几个要点步骤声明资源cpu_count2、gpu_count0注释明确Set to 1 for GPU training、memory4GB见 steps/train.pyEarlyStopping监视val_losspatience10防止无效训练浪费时间deterministicTrue保证结果可复现log_metadata(..., artifact_nametrial_result, infer_artifactTrue)把指标挂到输出构件上这是 Dashboard 可查询性的来源返回的 dict 同时携带n_epochs实际训练轮数与max_epochs上限。3. 汇总结果steps/report.py汇总所有试验结果并找出最优配置。跨轮结果通过 ZenML 构件previous_summary累积而不是共享数据库step def report_results(results: list, previous_summary: dict None) - dict: all_trials previous_summary.get(all_trials, []) current_results best min(all_trials, keylambda t: t[val_loss]) return { best_trial_number: best[trial_number], best_val_loss: best[val_loss], best_params: {...}, n_trials: len(all_trials), all_trials: all_trials, }该步骤以val_loss最小为最优判据与 study 的directionminimize一致返回的 summary 同时是下一轮suggest_trials的记忆载体。4. 重训最优模型steps/save_best.py找到最优超参数后用更多 epoch 重新训练以投入生产step def retrain_best_model(sweep_summary: dict, max_iter: int 20) - model: best_params sweep_summary[best_params] model FashionMNISTClassifier( learning_ratebest_params[learning_rate], hidden_dimbest_params[hidden_dim], ) trainer Trainer( max_epochsmax_iter, callbacks[EarlyStopping(monitorval_loss, patience15, modemin)], ... ) trainer.fit(model, train_loader, val_loader) log_metadata(metadata{...}, infer_artifactTrue) return model # 作为 ZenML 构件保存源码细节生产训练的EarlyStopping.patience提高到 15试验阶段为 10并关闭 checkpointing管线调用时把max_iter显式设为 200见 pipelines/sweep.py即试验只跑 10 epoch、生产模型跑 200 epoch返回的LightningModule被保存为 ZenML 构件best_model元数据中额外记录converged: trainer.current_epoch 1 max_iter用于判断早停是否触发并将hyperparameters、metrics、training分组写入元数据便于检索。为什么采用这种设计无存储浪费试验模型丢弃只保留指标生产就绪最优模型获得更长训练200 vs 10 epoch可查询模型元数据可被 Dashboard/MCP 查询可复现超参数全部记录随时可重训。5. 动态管线编排pipelines/sweep.py两条管线都用pipeline(dynamicTrue)声明通过.map()扇出试验pipeline(dynamicTrue) def sweep_pipeline( study_name: str fashion_mnist_sweep, n_trials: int 5, max_iter: int 10, ) - None: trials suggest_trials(study_namestudy_name, n_trialsn_trials) results train_trial.with_options(parameters{max_iter: max_iter}).map( trial_configtrials ) summary report_results(resultsresults) _ retrain_best_model(sweep_summarysummary, max_iter200)而自适应版本把上述过程放进for round_idx in range(n_rounds)循环每轮把上一轮的summary作为previous_summary传给suggest_trials实现边跑边学pipeline(dynamicTrue) def adaptive_sweep_pipeline( study_name: str fashion_mnist_sweep, n_rounds: int 3, trials_per_round: int 3, max_iter: int 10, ) - None: summary {round: 0, is_initial: True} for round_idx in range(n_rounds): trials suggest_trials( study_namestudy_name, n_trialstrials_per_round, previous_summarysummary, ) results train_trial.with_options( parameters{max_iter: max_iter} ).map(trial_configtrials) summary report_results(resultsresults, previous_summarysummary) _ retrain_best_model(sweep_summarysummary, max_iter200)dynamicTrue是 ZenML 动态管线dynamic pipelines的声明入口对应框架层实现见 src/zenml/pipelines/pipeline_decorator.pydynamic: Optional[bool] None与 src/zenml/pipelines/dynamic/pipeline_definition.py。管线运行快照会携带is_dynamic标志参与执行决策见 src/zenml/pipelines/build_utils.py。ZenML 自动替你处理了并行执行所有试验同时运行或按资源池排队构件追踪每个试验的输入/输出都被版本化缓存相同参数重复运行直接复用缓存结果元数据每个试验的指标可查询。两种优化模式对比模式一简单并行扫描Simple Sweepsweep_pipeline(n_trials10)一次性建议全部 10 个试验10 个试验并行训练或按资源可用性排队适用场景网格搜索、小规模随机搜索、快速验证实验。模式二自适应多轮扫描Adaptive Sweepadaptive_sweep_pipeline(n_rounds5, trials_per_round4)第 1 轮建议 4 个试验并训练、汇总第 2 轮Optuna 利用第 1 轮结果给出更优建议第 3 轮及以后基于全部历史持续精化适用场景大规模搜索空间、贝叶斯优化、追求样本效率。自定义搜索空间编辑 steps/suggest.py 中的SEARCH_SPACE常量SEARCH_SPACE: Dict[str, Any] { learning_rate: {low: 1e-4, high: 1e-2, log: True}, batch_size: {choices: [64, 128]}, hidden_dim: {choices: [8, 16, 32]}, }三个键分别对应learning_rate对数均匀浮点区间、batch_size类别选择、hidden_dim类别选择控制 CNN 首层卷积核数量。_build_distributions会自动把该字典翻译为 Optuna 分布对象新增超参数维度时需同步扩展该函数与suggest_trials中的suggest_*调用。远程执行与生产部署可移植架构本示例使用从 ZenML 构件历史播种的内存 Optuna study意味着容器之间无需共享数据库可运行在任意编排器上本地、Kubernetes、云端跨轮试验历史通过 ZenML 构件保留。资源池Resource Pools在 ZenML 中配置资源池后管线会自动跨可用资源管理并行执行无需改代码——同一个train_trial步骤在本地与生产环境行为一致。每个步骤声明的ResourceSettingsCPU、GPU、内存为调度器提供约束依据。并行下载的数据安全细节一个值得注意的工程细节多个并行试验可能同时触发 FashionMNIST 数据集下载steps/data.py 使用filelock.FileLock保护下载目录防止并发下载破坏缓存数据以固定种子seed42做 80/20 训练/验证划分保证可复现性。项目结构速览examples/optuna_hyperparameter_tuning/ -- README.md # 本文对应的官方说明 -- requirements.txt # 依赖optuna, torch, lightning 等 -- run.py # 入口simple/adaptive 两种模式 -- config/ | -- simple.yaml # sweep_pipeline 配置 | -- adaptive.yaml # adaptive_sweep_pipeline 配置 -- pipelines/ | -- __init__.py | -- sweep.py # sweep_pipeline adaptive_sweep_pipeline -- steps/ -- __init__.py -- model.py # 共享 FashionMNISTClassifier 定义 -- data.py # 共享 FashionMNIST 数据加载 -- suggest.py # Optuna ask API 试验建议 -- train.py # 单个试验训练PyTorch Lightning -- save_best.py # 重训最优模型用于生产 -- report.py # 结果聚合 汇总排障CategoricalDistribution does not support dynamic value space如果你在运行之间修改了搜索空间的类别取值就可能触发此错误。原因在于 study 是内存态、由构件历史播种历史构件仍带有旧的类别取值而新 study 的分布却基于新的类别集合两者冲突。解决办法修改SEARCH_SPACE后在配置文件config/simple.yaml或config/adaptive.yaml中更换study_name开启一个全新的 study即可避开旧构件历史。延伸阅读ZenML 动态管线dynamic pipelines框架实现src/zenml/pipelines/dynamic/pipeline_definition.py、src/zenml/pipelines/pipeline_decorator.pyZenML 管线配置与 YAML 用法docs/book/how-to/steps-pipelines/目录本文所有示例代码均可直接在 examples/optuna_hyperparameter_tuning 中查看与复现。【免费下载链接】zenmlZenML : One AI Platform from Pipelines to Agents. https://zenml.io.项目地址: https://gitcode.com/GitHub_Trending/ze/zenml创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考