mkurman/optuna

---|----------|-------------| | TPESampler | General ML tuning | Tree-structured Parzen Estimator; default, good for most cases | | CMAESSampler | Continuous, low-dim (<100) | Covariance Matrix Adaptation; efficient for numeric params | | NSGAIISampler | Multi-objective (2-3 objectives) | Pareto-front optimization | | GPSampler | Expensive evaluations | Gaussian Process-based; sample-efficient | | RandomSampler | Baseline, debugging | Uniform random sampling | | GridSampler | Small discrete spaces | Exhaustive grid search | | QMCSampler | Continuous spaces | Quasi-Monte Carlo, better coverage than random |

Usage:

python
import optuna
sampler = optuna.samplers.TPESampler(seed=42, n_startup_trials=10)
study = optuna.create_study(sampler=sampler)

4. Pruning (Early Stopping)

Stop unpromising trials early to save compute:

python
def objective(trial):
    for epoch in range(100):
        accuracy = train_and_evaluate(...)
        # Report intermediate value
        trial.report(accuracy, epoch)
        # Check if should prune
        if trial.should_prune():
            raise optuna.TrialPruned()
    return accuracy

Pruner Selection:

• MedianPruner: Prune if trial's intermediate value is below median at same step
• HyperbandPruner: Successive halving; efficient for large trial counts
• SuccessiveHalvingPruner: Similar to Hyperband, simpler configuration
• ThresholdPruner: Prune below absolute threshold
• PatientPruner: Prune after N epochs without improvement

Integration with PyTorch Lightning:

python
from optuna.integration import PyTorchLightningPruningCallback

trainer = pl.Trainer(
    callbacks=[PyTorchLightningPruningCallback(trial, monitor="val_acc")],
    max_epochs=100,
)

5. Multi-Objective Optimization

python
def objective(trial):
    accuracy = train_and_get_accuracy(trial)
    latency_ms = measure_latency(trial)
    return accuracy, latency_ms  # Return tuple

study = optuna.create_study(
    directions=["maximize", "minimize"],
    sampler=optuna.samplers.NSGAIISampler(),
)

study.optimize(objective, n_trials=200)

# Get Pareto front
best_trials = study.best_trials
for trial in best_trials:
    print(f"Params: {trial.params}, Values: {trial.values}")

6. Distributed / Parallel Execution

Single-machine multi-process:

python
study.optimize(objective, n_trials=100, n_jobs=8)  # 8 parallel workers

Multi-node via shared storage (SQLite):

python
# On all nodes, share the same study name and storage
study = optuna.create_study(
    study_name="distributed_study",
    storage="sqlite:///optuna_study.db",
    load_if_exists=True,
)
study.optimize(objective, n_trials=500)

Multi-node via RDB (PostgreSQL/MySQL):

python
study = optuna.create_study(
    study_name="large_scale_study",
    storage="postgresql://user:pass@host:5432/optuna",
    load_if_exists=True,
)

7. Visualization

python
from optuna.visualization import (
    plot_optimization_history,
    plot_param_importances,
    plot_parallel_coordinate,
    plot_contour,
    plot_slice,
)

# Optimization progress over trials
plot_optimization_history(study)

# Hyperparameter importance ranking
plot_param_importances(study)

# Parallel coordinate plot for high-dimensional analysis
plot_parallel_coordinate(study)

# Slice plot showing parameter-value relationship
plot_slice(study)

# Contour plot for pairwise parameter interactions
plot_contour(study, params=["learning_rate", "n_layers"])

Web Dashboard (optuna-dashboard):

bash
pip install optuna-dashboard
optuna-dashboard sqlite:///optuna_study.db
# Opens at http://localhost:8080

8. PyTorch Lightning Integration

python
import pytorch_lightning as pl
import optuna
from optuna.integration import PyTorchLightningPruningCallback

def objective(trial):
    # Suggest hyperparameters
    lr = trial.suggest_float("lr", 1e-5, 1e-1, log=True)
    batch_size = trial.suggest_categorical("batch_size", [32, 64, 128, 256])
    n_layers = trial.suggest_int("n_layers", 1, 6)

    model = MyLightningModule(lr=lr, n_layers=n_layers)
    trainer = pl.Trainer(
        max_epochs=50,
        callbacks=[PyTorchLightningPruningCallback(trial, monitor="val_loss")],
        logger=False,
    )
    trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)

    return trainer.callback_metrics["val_loss"].item()

study = optuna.create_study(direction="minimize")
study.optimize(objective, n_trials=50)

9. HuggingFace Transformers Integration

python
from transformers import Trainer, TrainingArguments
import optuna

def hp_space(trial):
    return {
        "learning_rate": trial.suggest_float("learning_rate", 1e-6, 1e-4, log=True),
        "per_device_train_batch_size": trial.suggest_categorical("batch_size", [8, 16, 32]),
        "num_train_epochs": trial.suggest_int("num_epochs", 1, 5),
        "warmup_ratio": trial.suggest_float("warmup_ratio", 0.0, 0.3),
    }

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

best_run = trainer.hyperparameter_search(
    hp_space=hp_space,
    n_trials=30,
    direction="minimize",
)

10. Artifact and Attribute Storage

python
def objective(trial):
    model = train_model(trial)
    # Store arbitrary attributes
    trial.set_user_attr("model_architecture", str(model))
    trial.set_user_attr("training_time_seconds", 3600)
    return evaluate(model)

# Retrieve later
for trial in study.trials:
    print(trial.user_attrs.get("training_time_seconds"))

Installation

bash
pip install optuna
# Optional: dashboard
pip install optuna-dashboard
# Optional: OptunaHub features
pip install optunahub

Key Patterns for ML Training

1. Always use log=True for learning rates, batch sizes, and other scale-sensitive params
2. Set n_startup_trials to 10-20 for TPE to warm up with random exploration
3. Use pruning aggressively for expensive deep learning trials — saves 50-80% compute
4. For reproducibility, set seed on both sampler and study.optimize()
5. Store intermediate values with trial.report() even if not pruning — enables better analysis

References

• Optuna Documentation
• Optuna Examples
• OptunaHub — community-contributed samplers, pruners, and visualization
• Optuna Dashboard

See scripts/optuna_lightning_template.py for a complete PyTorch Lightning + Optuna training template. See references/advanced_samplers.md for detailed sampler comparison and selection guidance.