Example 14: Hyperparameter Optimization with Ray Tune¶

Manually tuning hyperparameters -- learning rate, hidden size, model type -- is tedious and error-prone. TSFast integrates with Ray Tune to automate the search. This example runs a small hyperparameter search to find the best model configuration for the Silverbox benchmark.

Prerequisites¶

This example builds on concepts from:

• Example 00 -- data loading and model training basics
• Example 04 -- model architectures and rnn_type

Make sure Ray Tune is installed:

uv sync --extra dev

Setup¶

In [1]:

from tsfast.tsdata.benchmark import create_dls_silverbox
from tsfast.models.rnn import RNNLearner
from tsfast.tune import HPOptimizer, log_uniform
from tsfast.training import fun_rmse
from ray import tune

from tsfast.tsdata.benchmark import create_dls_silverbox from tsfast.models.rnn import RNNLearner from tsfast.tune import HPOptimizer, log_uniform from tsfast.training import fun_rmse from ray import tune

/home/pheenix/Development/tsfast/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

2026-03-01 00:06:19,946	INFO util.py:154 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.

2026-03-01 00:06:20,190	INFO util.py:154 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.

Why Hyperparameter Optimization?¶

Model performance depends heavily on hyperparameters: learning rate, hidden size, architecture choice, and regularization strength. Finding the right combination by hand requires many experiments and careful record-keeping.

Automated approaches help:

• Grid search evaluates every combination -- thorough but expensive.
• Random search samples randomly and is surprisingly effective in high-dimensional spaces.
• Population-based training evolves configurations during training, combining exploration with exploitation.

Ray Tune provides all of these strategies (and more) behind a unified API. TSFast's HPOptimizer wraps Ray Tune so you can search over model configurations with minimal boilerplate.

Prepare the DataLoaders¶

We use the Silverbox benchmark with a small batch size and window size to keep the example lightweight.

In [2]:

dls = create_dls_silverbox(bs=16, win_sz=500, stp_sz=10)

dls = create_dls_silverbox(bs=16, win_sz=500, stp_sz=10)

Define a Learner Factory¶

HPOptimizer needs a factory function that takes (dls, config) and returns a configured Learner. Ray Tune calls this function once per trial, each time with a different hyperparameter configuration sampled from the search space.

In [3]:

def create_learner(dls, config):
    """Create a configured RNNLearner from hyperparameter config."""
    return RNNLearner(
        dls,
        rnn_type=config["rnn_type"],
        hidden_size=config["hidden_size"],
        n_skip=50,
        metrics=[fun_rmse],
    )

def create_learner(dls, config): """Create a configured RNNLearner from hyperparameter config.""" return RNNLearner( dls, rnn_type=config["rnn_type"], hidden_size=config["hidden_size"], n_skip=50, metrics=[fun_rmse], )

Define the Search Space¶

The search space is a plain dictionary where values are Ray Tune sampling primitives:

• tune.choice -- samples uniformly from a list of discrete options. Good for categorical parameters like architecture type or layer count.
• log_uniform -- samples uniformly on a logarithmic scale. Ideal for parameters that span orders of magnitude, such as learning rate.

We start with a small search over two parameters: RNN cell type and hidden size.

In [4]:

search_config = {
    "rnn_type": tune.choice(["gru", "lstm"]),
    "hidden_size": tune.choice([32, 40]),
    "n_epoch": 3,
    "lr": 3e-3,
}

search_config = { "rnn_type": tune.choice(["gru", "lstm"]), "hidden_size": tune.choice([32, 40]), "n_epoch": 3, "lr": 3e-3, }

The config also contains fixed training parameters:

• n_epoch=3 -- each trial trains for 3 epochs (enough to compare configurations, not enough for final training).
• lr=3e-3 -- fixed learning rate for all trials in this first search.

Run the Optimization¶

HPOptimizer takes the learner factory and the DataLoaders. Calling optimize launches the search: num_samples=4 runs 4 independent trials, each with a different hyperparameter combination drawn from search_config.

The default training function uses fit_flat_cos and reports training loss, validation loss, and metrics to Ray Tune after every epoch.

In [5]:

optimizer = HPOptimizer(
    create_lrn=create_learner,
    dls=dls,
)

results = optimizer.optimize(
    config=search_config,
    num_samples=4,
    resources_per_trial={"cpu": 1, "gpu": 0},
)

optimizer = HPOptimizer( create_lrn=create_learner, dls=dls, ) results = optimizer.optimize( config=search_config, num_samples=4, resources_per_trial={"cpu": 1, "gpu": 0}, )

Tune Status

Current time:	2026-03-01 00:08:10
Running for:	00:01:45.94
Memory:	11.4/15.6 GiB

System Info

Using FIFO scheduling algorithm.
Logical resource usage: 1.0/24 CPUs, 0/1 GPUs (0.0/1.0 accelerator_type:G)

Trial Status

Trial name	status	loc	hidden_size	rnn_type	iter	total time (s)	train_loss	valid_loss	fun_rmse
learner_optimize_1aa84_00000	TERMINATED	172.24.182.240:56506	40	gru	3	94.3783	0.0012801	0.00100963	0.00220393
learner_optimize_1aa84_00001	TERMINATED	172.24.182.240:56507	32	gru	3	90.8211	0.00133288	0.00105418	0.00219282
learner_optimize_1aa84_00002	TERMINATED	172.24.182.240:56508	40	gru	3	95.3377	0.00129515	0.000939281	0.00197045
learner_optimize_1aa84_00003	TERMINATED	172.24.182.240:56514	32	lstm	3	34.7632	0.00125725	0.000982948	0.00206519

(raylet) warning: `VIRTUAL_ENV=/home/pheenix/Development/tsfast/.venv` does not match the project environment path `.venv` and will be ignored; use `--active` to target the active environment instead
(raylet) Using CPython 3.12.12
(raylet) Creating virtual environment at: .venv

(raylet)    Building tsfast @ file:///tmp/ray/session_2026-03-01_00-06-20_700222_53928/runtime_resources/working_dir_files/_ray_pkg_e0e5a818bebb9c81

(raylet) Installed 165 packages in 1.00s

(learner_optimize pid=56506) /tmp/ray/session_2026-03-01_00-06-20_700222_53928/runtime_resources/working_dir_files/_ray_pkg_e0e5a818bebb9c81/.venv/lib/python3.12/site-packages/torch/utils/data/dataloader.py:775: UserWarning: 'pin_memory' argument is set as true but no accelerator is found, then device pinned memory won't be used.
(learner_optimize pid=56506)   super().__init__(loader)
(raylet) warning: `VIRTUAL_ENV=/home/pheenix/Development/tsfast/.venv` does not match the project environment path `.venv` and will be ignored; use `--active` to target the active environment instead [repeated 7x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/user-guides/configure-logging.html#log-deduplication for more options.)
(raylet)       Built tsfast @ file:///tmp/ray/session_2026-03-01_00-06-20_700222_53928/runtime_resources/working_dir_files/_ray_pkg_e0e5a818bebb9c81

(learner_optimize pid=56514) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/pheenix/ray_results/learner_optimize_2026-03-01_00-06-24/learner_optimize_1aa84_00003_3_hidden_size=32,rnn_type=lstm_2026-03-01_00-06-24/checkpoint_000000)
(learner_optimize pid=56514) /tmp/ray/session_2026-03-01_00-06-20_700222_53928/runtime_resources/working_dir_files/_ray_pkg_e0e5a818bebb9c81/.venv/lib/python3.12/site-packages/torch/utils/data/dataloader.py:775: UserWarning: 'pin_memory' argument is set as true but no accelerator is found, then device pinned memory won't be used. [repeated 7x across cluster]
(learner_optimize pid=56514)   super().__init__(loader) [repeated 7x across cluster]

(learner_optimize pid=56514) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/pheenix/ray_results/learner_optimize_2026-03-01_00-06-24/learner_optimize_1aa84_00003_3_hidden_size=32,rnn_type=lstm_2026-03-01_00-06-24/checkpoint_000001)

(learner_optimize pid=56507) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/pheenix/ray_results/learner_optimize_2026-03-01_00-06-24/learner_optimize_1aa84_00001_1_hidden_size=32,rnn_type=gru_2026-03-01_00-06-24/checkpoint_000000)

(learner_optimize pid=56507) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/pheenix/ray_results/learner_optimize_2026-03-01_00-06-24/learner_optimize_1aa84_00001_1_hidden_size=32,rnn_type=gru_2026-03-01_00-06-24/checkpoint_000001) [repeated 4x across cluster]

(learner_optimize pid=56507) Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/pheenix/ray_results/learner_optimize_2026-03-01_00-06-24/learner_optimize_1aa84_00001_1_hidden_size=32,rnn_type=gru_2026-03-01_00-06-24/checkpoint_000002) [repeated 3x across cluster]

2026-03-01 00:08:10,643	INFO tune.py:1009 -- Wrote the latest version of all result files and experiment state to '/home/pheenix/ray_results/learner_optimize_2026-03-01_00-06-24' in 0.0049s.

2026-03-01 00:08:10,648	INFO tune.py:1041 -- Total run time: 105.98 seconds (105.93 seconds for the tuning loop).

Analyze Results¶

The optimize call returns a Ray Tune ExperimentAnalysis object stored in optimizer.analysis. You can query it for the best trial configuration, inspect per-trial results, or export data for further analysis.

In [6]:

best = optimizer.analysis.get_best_config(metric="valid_loss", mode="min")
print("Best config:")
for key in ["rnn_type", "hidden_size", "lr"]:
    print(f"  {key}: {best[key]}")

best = optimizer.analysis.get_best_config(metric="valid_loss", mode="min") print("Best config:") for key in ["rnn_type", "hidden_size", "lr"]: print(f" {key}: {best[key]}")

Best config:
  rnn_type: gru
  hidden_size: 40
  lr: 0.003

In [7]:

result_df = optimizer.analysis.results_df
print("\nAll trial results:")
result_df[["config/rnn_type", "config/hidden_size", "valid_loss"]]

result_df = optimizer.analysis.results_df print("\nAll trial results:") result_df[["config/rnn_type", "config/hidden_size", "valid_loss"]]

All trial results:

Out[7]:

	config/rnn_type	config/hidden_size	valid_loss
trial_id
1aa84_00000	gru	40	0.001010
1aa84_00001	gru	32	0.001054
1aa84_00002	gru	40	0.000939
1aa84_00003	lstm	32	0.000983

Using log_uniform for Learning Rate¶

In the first search we fixed the learning rate. A more thorough search treats lr as a tunable parameter using log_uniform. This samples on a logarithmic scale between the given bounds -- appropriate because the difference between 1e-4 and 1e-3 matters more than between 1e-2 and 1.1e-2.

In [8]:

search_config_v2 = {
    "rnn_type": tune.choice(["gru", "lstm"]),
    "hidden_size": tune.choice([32, 40]),
    "lr": log_uniform(1e-4, 1e-2),
    "n_epoch": 3,
}

search_config_v2 = { "rnn_type": tune.choice(["gru", "lstm"]), "hidden_size": tune.choice([32, 40]), "lr": log_uniform(1e-4, 1e-2), "n_epoch": 3, }

When lr is a callable sampler in the config, the training function samples a fresh value for each trial. This overrides any fixed learning rate.

In [9]:

optimizer_v2 = HPOptimizer(
    create_lrn=create_learner,
    dls=dls,
)

results_v2 = optimizer_v2.optimize(
    config=search_config_v2,
    num_samples=4,
    resources_per_trial={"cpu": 1, "gpu": 0},
)

optimizer_v2 = HPOptimizer( create_lrn=create_learner, dls=dls, ) results_v2 = optimizer_v2.optimize( config=search_config_v2, num_samples=4, resources_per_trial={"cpu": 1, "gpu": 0}, )

Tune Status

Current time:	2026-03-01 00:09:46
Running for:	00:01:36.08
Memory:	8.8/15.6 GiB

System Info

Using FIFO scheduling algorithm.
Logical resource usage: 1.0/24 CPUs, 0/1 GPUs (0.0/1.0 accelerator_type:G)

Trial Status

Trial name	status	loc	hidden_size	rnn_type	iter	total time (s)	train_loss	valid_loss	fun_rmse
learner_optimize_59dd8_00000	TERMINATED	172.24.182.240:62361	40	lstm	3	27.1214	0.00134441	0.000934431	0.0019721
learner_optimize_59dd8_00001	TERMINATED	172.24.182.240:62353	32	lstm	3	26.3676	0.00132948	0.00119851	0.00227271
learner_optimize_59dd8_00002	TERMINATED	172.24.182.240:62362	40	gru	3	92.3884	0.00140121	0.00130199	0.00259187
learner_optimize_59dd8_00003	TERMINATED	172.24.182.240:62363	40	gru	3	89.0636	0.00155123	0.00112293	0.00229254

2026-03-01 00:09:46,816	INFO tune.py:1009 -- Wrote the latest version of all result files and experiment state to '/home/pheenix/ray_results/learner_optimize_2026-03-01_00-08-10' in 0.0123s.

2026-03-01 00:09:46,820	INFO tune.py:1041 -- Total run time: 96.10 seconds (96.07 seconds for the tuning loop).

In [10]:

best_v2 = optimizer_v2.analysis.get_best_config(metric="valid_loss", mode="min")
print("Best config (with lr search):")
for key in ["rnn_type", "hidden_size", "lr"]:
    print(f"  {key}: {best_v2[key]}")

best_v2 = optimizer_v2.analysis.get_best_config(metric="valid_loss", mode="min") print("Best config (with lr search):") for key in ["rnn_type", "hidden_size", "lr"]: print(f" {key}: {best_v2[key]}")

Best config (with lr search):
  rnn_type: lstm
  hidden_size: 40
  lr: <function log_uniform.<locals>._sample at 0x7dc448ff4900>

Key Takeaways¶

• HPOptimizer wraps Ray Tune for easy hyperparameter search with TSFast. Pass a learner factory and DataLoaders, then call optimize.
• Learner factory -- a function (dls, config) -> Learner that builds a fresh model from the hyperparameter config each trial.
• tune.choice for categorical parameters (architecture, layer count); log_uniform for continuous parameters on a log scale (learning rate).
• Start small -- few trials, few epochs -- to validate the pipeline before scaling up.
• optimizer.analysis gives access to the full Ray Tune ExperimentAnalysis for querying best configs, exporting results, and loading the best checkpoint.