DeepAdvisors | Hyperparameter Tuning TensorFlow Models with KerasTuner: Search Spaces, Hyperband, and Distributed Sweeps

Every client engagement that starts with "the model is not accurate enough" eventually reaches the same question: have you actually searched the hyperparameters, or did you hand-tune three settings in a notebook and stop? Hand-tuning is fine for a first baseline. It is a poor way to spend the last few points of accuracy, and it is unreproducible - nobody can tell you six months later why the learning rate is 3e-4.

This tutorial covers doing it properly on the TensorFlow/Keras stack with KerasTuner: defining a search space, choosing a search algorithm, running the sweep across several GPUs or machines, and turning the result into something you can put in a pipeline.

Install

pip install -U keras keras-tuner tensorboard

KerasTuner works with Keras 3 on any backend, so a sweep you write today runs on TensorFlow now and JAX later. The API reference lives at keras.io/keras_tuner.

Step 1: make the model a function of hyperparameters

The core idea is a build function that takes an hp object and returns a compiled model. Anything you sample from hp becomes part of the search space.

import keras
import keras_tuner as kt

NUM_FEATURES = 64
NUM_CLASSES = 4


def build_model(hp):
    model = keras.Sequential()
    model.add(keras.layers.Input(shape=(NUM_FEATURES,)))

    for i in range(hp.Int("num_layers", 1, 4)):
        model.add(
            keras.layers.Dense(
                units=hp.Int(f"units_{i}", min_value=32, max_value=512, step=32),
                activation=hp.Choice("activation", ["relu", "gelu"]),
            )
        )
        if hp.Boolean("batch_norm"):
            model.add(keras.layers.BatchNormalization())
        model.add(keras.layers.Dropout(hp.Float("dropout", 0.0, 0.5, step=0.1)))

    model.add(keras.layers.Dense(NUM_CLASSES, activation="softmax"))

    model.compile(
        optimizer=keras.optimizers.AdamW(
            learning_rate=hp.Float("lr", 1e-4, 1e-2, sampling="log"),
            weight_decay=hp.Float("weight_decay", 1e-6, 1e-2, sampling="log"),
        ),
        loss="sparse_categorical_crossentropy",
        metrics=["sparse_categorical_accuracy"],
    )
    return model

Three things worth copying from that snippet:

sampling="log" for learning rate and weight decay. Learning rate matters on a multiplicative scale; a linear sweep from 1e-4 to 1e-2 spends almost all of its budget on large values.
Conditional depth. units_{i} inside the num_layers loop means a one-layer trial never samples hyperparameters it does not use. KerasTuner handles that conditionality for you.
Keep the space small. Every extra dimension costs trials. Start with learning rate, width, depth, and regularisation. Add exotic knobs only after those stop paying.

Step 2: pick a tuner

KerasTuner ships several search strategies. In practice:

Tuner	Use when
`RandomSearch`	Baseline and sanity check. Surprisingly strong; beats a grid of the same size.
`Hyperband`	Default choice for deep nets. Trains many configs briefly, promotes survivors. Best accuracy per GPU-hour.
`BayesianOptimization`	Expensive trials (hours each) and a small, mostly continuous space.
`GridSearch`	Only for tiny, genuinely categorical spaces you must report exhaustively.

tuner = kt.Hyperband(
    build_model,
    objective=kt.Objective("val_sparse_categorical_accuracy", direction="max"),
    max_epochs=40,
    factor=3,
    hyperband_iterations=2,
    directory="tuning",
    project_name="churn_mlp",
    seed=1337,
)
tuner.search_space_summary()

Name the objective explicitly. The most common self-inflicted wound in a sweep is tuning against training accuracy, or against a custom metric whose direction the tuner has to guess. If you optimise a custom metric, wrap it in kt.Objective and state the direction.

Choose an objective that matches the business decision, too. For an imbalanced classifier, validation AUC or PR-AUC beats accuracy; for a ranking model, use the retrieval metric you will actually be judged on in production.

Step 3: run the search

search() takes the same arguments as fit():

callbacks = [
    keras.callbacks.EarlyStopping(
        monitor="val_sparse_categorical_accuracy",
        patience=5,
        restore_best_weights=True,
    ),
    keras.callbacks.TensorBoard("tuning/tb"),
]

tuner.search(
    train_ds,            # a tf.data.Dataset
    validation_data=val_ds,
    epochs=40,
    callbacks=callbacks,
    verbose=2,
)

If your input pipeline is the bottleneck, the sweep multiplies that cost by the number of trials - cache and prefetch before you start, as in our tf.data tutorial. A sweep is the fastest way to discover that your GPU was idle 70% of the time all along.

You can tune fit-time hyperparameters as well by subclassing HyperModel:

class ChurnHyperModel(kt.HyperModel):
    def build(self, hp):
        return build_model(hp)

    def fit(self, hp, model, *args, **kwargs):
        return model.fit(
            *args,
            batch_size=hp.Choice("batch_size", [128, 256, 512]),
            **kwargs,
        )

That is also where augmentation strength, label smoothing, or a class-sampling ratio belong - anything that is part of the training loop rather than the architecture.

Step 4: read the results properly

tuner.results_summary(num_trials=10)

best_hp = tuner.get_best_hyperparameters(num_trials=1)[0]
print(best_hp.values)

Do not ship tuner.get_best_models()[0]. Under Hyperband the best trial may have been trained on only a fraction of the epoch budget, and its weights were selected on the validation split. Retrain the winning configuration from scratch, then score it once on a held-out test set:

final = build_model(best_hp)
final.fit(full_train_ds, epochs=best_epochs, callbacks=[...])
final.evaluate(test_ds)
final.save("artifacts/churn_mlp.keras")

Look past the single best trial as well. If the top ten trials all sit at the edge of your lr range, widen the range and rerun. If accuracy varies more between repeated runs of the same configuration than between the top configurations, your sweep is measuring seed noise - fix the seed, average over three seeds per configuration, or accept the simpler model.

Step 5: distribute the sweep

There are two independent axes of parallelism.

Multiple GPUs per trial. Train each trial data-parallel by handing the tuner a distribution strategy, so one trial uses the whole machine:

import tensorflow as tf

tuner = kt.Hyperband(
    build_model,
    objective="val_sparse_categorical_accuracy",
    max_epochs=40,
    distribution_strategy=tf.distribute.MirroredStrategy(),
    directory="tuning",
    project_name="churn_mlp_dp",
)

See scaling training beyond one GPU for the trade-offs between this and larger batch sizes.

Multiple trials in parallel. KerasTuner has a chief/worker mode driven entirely by environment variables plus a shared directory on network storage:

# chief
export KERASTUNER_TUNER_ID="chief"
export KERASTUNER_ORACLE_IP="10.0.0.5"
export KERASTUNER_ORACLE_PORT="8000"
python tune.py &

# each worker
export KERASTUNER_TUNER_ID="tuner0"   # tuner1, tuner2, ...
export KERASTUNER_ORACLE_IP="10.0.0.5"
export KERASTUNER_ORACLE_PORT="8000"
python tune.py

Every process runs the same script; the oracle hands out trials. This is the cheapest way to fill a fleet of spot GPU instances, because a worker dying only loses its current trial - the oracle state lives with the chief and the shared directory.

Step 6: put the sweep in the pipeline, not in a notebook

Once you trust the search, the artefact you keep is a hyperparameter JSON, not a model:

import json

with open("artifacts/best_hp.json", "w") as f:
    json.dump(best_hp.values, f, indent=2)

Commit that file. Your retraining job reads it and trains deterministically, which is what makes a TFX continuous-training pipeline reproducible - TFX even has a Tuner component if you want the search itself inside the DAG. Re-run the search on a schedule (quarterly, or when TensorFlow Data Validation flags a real distribution shift) rather than on every training run, which would make every retrain a different model.

Budgeting a sweep

A rough rule we use when scoping client work: a Hyperband sweep with factor=3 costs a small multiple of training one model for the full epoch budget, per Hyperband iteration, while evaluating dozens of configurations. For a model that will run in production for a year, that is an excellent trade. What it is not good for is a model whose ceiling is set by features or labels - if error analysis says 8% of your labels are wrong, no hyperparameter will save you, and the tuning budget belongs in the data.

Common failure modes

Validation split too small. With 500 validation examples the top 20 trials are statistically indistinguishable and you will overfit the split. Use a larger split or cross-validated trials.
Leaky preprocessing. Adapt Normalization layers and vocabularies on training data inside the trial. Adapting on the full dataset before the sweep leaks validation statistics into every trial.
Resuming into a stale directory. KerasTuner resumes from directory/project_name. That is a feature after a spot interruption and a bug once you have changed the search space - bump project_name.
No early stopping. Without it, hopeless configurations burn the full epoch budget.
Tuning before profiling. If one trial takes four hours because of a slow input pipeline, fix that first; you get ten times the trials for the same money.

Want this done for you?

We run hyperparameter searches, distributed training, and reproducible retraining pipelines as a standard part of our TensorFlow consulting engagements - usually a short, fixed-scope sprint that ends with a committed best_hp.json, a retraining script, and an honest statement of how much headroom is left in the model. If that is the shape of your problem, get in touch.