Hyperparameter tuning

Running a hyperparameter tuning trial with ZenML.

Hyperparameter tuning is not yet a first-class citizen in ZenML, but it is (high up) on our roadmap of features and will likely receive first-class ZenML support soon. In the meanwhile, the following example shows how hyperparameter tuning can currently be implemented within a ZenML run.

A basic iteration through a number of hyperparameters can be achieved with ZenML by using a simple pipeline like this:

@pipeline

def my_pipeline(step_count: int) -> None:

data = load_data_step()

after = []

for i in range(step_count):

train_step(data, learning_rate=i * 0.0001, name=f"train_step_{i}")

after.append(f"train_step_{i}")

model = select_model_step(..., after=after)

This is an implementation of a basic grid search (across a single dimension) that would allow for a different learning rate to be used across the same train_step. Once that step has been run for all the different learning rates, the select_model_step finds which hyperparameters gave the best results or performance.

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

pipelines/training.py, you will find a training pipeline with a

...

########## Hyperparameter tuning stage ##########

after = []

search_steps_prefix = "hp_tuning_search_"

for i, model_search_configuration in enumerate(

MetaConfig.model_search_space

):

step_name = f"{search_steps_prefix}{i}"

hp_tuning_single_search(

model_metadata=ExternalArtifact(

value=model_search_configuration,

),

id=step_name,

dataset_trn=dataset_trn,

dataset_tst=dataset_tst,

target=target,

after.append(step_name)

best_model_config = hp_tuning_select_best_model(

search_steps_prefix=search_steps_prefix, after=after

...