steps_prefix=search_steps_prefix, after=after

...The main challenge of this implementation is that it is currently not possible to pass a variable number of artifacts into a step programmatically, so the select_model_step needs to query all artifacts produced by the previous steps via the ZenML Client instead:

from zenml import step, get_step_context

from zenml.client import Client

@step

def select_model_step():

run_name = get_step_context().pipeline_run.name

run = Client().get_pipeline_run(run_name)

# Fetch all models trained by a 'train_step' before

trained_models_by_lr = {}

for step_name, step in run.steps.items():

if step_name.startswith("train_step"):

for output_name, output in step.outputs.items():

if output_name == "<NAME_OF_MODEL_OUTPUT_IN_TRAIN_STEP>":

model = output.load()

lr = step.config.parameters["learning_rate"]

trained_models_by_lr[lr] = model

# Evaluate the models to find the best one

for lr, model in trained_models_by_lr.items():

...

To set up the local environment used below, follow the recommendations from the Project templates.

In the steps/hp_tuning folder, you will find two step files, which can be used as a starting point for building your own hyperparameter search tailored specifically to your use case:

hp_tuning_single_search(...) is performing a randomized search for the best model hyperparameters in a configured space.

hp_tuning_select_best_model(...) is searching for the best hyperparameters, looping other results of previous random searches to find the best model according to a defined metric.

PreviousUse failure/success hooks

NextVersion pipelines

Last updated 15 days ago