gavindiaz/rust-optimizer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement Parameters API

Browse Source
This commit is contained in:
Manuel Raimann
2026-02-06 17:15:30 +01:00
parent 2fef1ab3c3
commit 55e50e6afb
19 changed files with 3023 additions and 2089 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/async_api_optimization.rs
+75 -60
View File
@@ -26,6 +26,7 @@
use std::time::{Duration, Instant};
use optimizer::parameter::{BoolParam, CategoricalParam, IntParam, Parameter};
use optimizer::sampler::tpe::TpeSampler;
use optimizer::{Direction, ParamValue, Study, Trial};
@@ -133,32 +134,27 @@ async fn evaluate_service(config: &ServiceConfig) -> f64 {
/// 3. Return both the Trial and the result value as a tuple
///
/// This ownership pattern allows the trial to be used across await points.
async fn objective(mut trial: Trial) -> optimizer::Result<(Trial, f64)> {
// Sample configuration parameters
// Stepped integers: only sample multiples of the step value
let cache_size_mb = trial.suggest_int_step("cache_size_mb", 64, 1024, 64)?;
let connection_pool_size = trial.suggest_int_step("connection_pool_size", 10, 200, 10)?;
let request_timeout_ms = trial.suggest_int_step("request_timeout_ms", 1000, 10000, 500)?;
// Regular integer
let retry_count = trial.suggest_int("retry_count", 0, 5)?;
// Log-scale integer: good for parameters like batch sizes
// that might vary from 1 to 256
let batch_size = trial.suggest_int_log("batch_size", 1, 256)?;
// Regular integer for compression level
let compression_level = trial.suggest_int("compression_level", 0, 9)?;
// Boolean: internally uses categorical with [false, true]
let use_http2 = trial.suggest_bool("use_http2")?;
// Categorical: choose from a list of options
let load_balancing = trial.suggest_categorical(
"load_balancing",
&["round_robin", "least_connections", "random", "ip_hash"],
)?;
#[allow(clippy::too_many_arguments)]
async fn objective(
mut trial: Trial,
cache_size_mb_param: &IntParam,
connection_pool_size_param: &IntParam,
request_timeout_ms_param: &IntParam,
retry_count_param: &IntParam,
batch_size_param: &IntParam,
compression_level_param: &IntParam,
use_http2_param: &BoolParam,
load_balancing_param: &CategoricalParam<&str>,
) -> optimizer::Result<(Trial, f64)> {
// Sample configuration parameters using parameter definitions
let cache_size_mb = cache_size_mb_param.suggest(&mut trial)?;
let connection_pool_size = connection_pool_size_param.suggest(&mut trial)?;
let request_timeout_ms = request_timeout_ms_param.suggest(&mut trial)?;
let retry_count = retry_count_param.suggest(&mut trial)?;
let batch_size = batch_size_param.suggest(&mut trial)?;
let compression_level = compression_level_param.suggest(&mut trial)?;
let use_http2 = use_http2_param.suggest(&mut trial)?;
let load_balancing = load_balancing_param.suggest(&mut trial)?;
// Build configuration
let config = ServiceConfig {
@@ -184,20 +180,11 @@ async fn objective(mut trial: Trial) -> optimizer::Result<(Trial, f64)> {
// ============================================================================
/// Formats a parameter value for display.
fn format_param(name: &str, value: &ParamValue) -> String {
match (name, value) {
(_, ParamValue::Float(v)) => format!("{v:.4}"),
(_, ParamValue::Int(v)) => format!("{v}"),
("use_http2", ParamValue::Categorical(idx)) => {
if *idx == 1 { "true" } else { "false" }.to_string()
}
("load_balancing", ParamValue::Categorical(idx)) => {
["round_robin", "least_connections", "random", "ip_hash"]
.get(*idx)
.unwrap_or(&"unknown")
.to_string()
}
(_, ParamValue::Categorical(idx)) => format!("category_{idx}"),
fn format_param(value: &ParamValue) -> String {
match value {
ParamValue::Float(v) => format!("{v:.4}"),
ParamValue::Int(v) => format!("{v}"),
ParamValue::Categorical(idx) => format!("category_{idx}"),
}
}
@@ -226,23 +213,13 @@ fn print_best_config(study: &Study<f64>) -> optimizer::Result<()> {
println!(" Score: {:.6}", best.value);
println!("\n Parameters:");
// Print parameters in a logical order
let param_order = [
"cache_size_mb",
"connection_pool_size",
"request_timeout_ms",
"retry_count",
"batch_size",
"compression_level",
"use_http2",
"load_balancing",
];
for name in param_order {
if let Some(value) = best.params.get(name) {
let display = format_param(name, value);
println!(" {name}: {display}");
}
for (id, value) in &best.params {
let label = best
.param_labels
.get(id)
.map_or_else(|| format!("{id}"), |l| l.clone());
let display = format_param(value);
println!(" {label}: {display}");
}
Ok(())
@@ -284,7 +261,22 @@ async fn main() -> optimizer::Result<()> {
// Step 2: Create a study to minimize the score
let study: Study<f64> = Study::with_sampler(Direction::Minimize, sampler);
// Step 3: Configure optimization
// Step 3: Define parameter search spaces
let cache_size_mb_param = IntParam::new(64, 1024).step(64);
let connection_pool_size_param = IntParam::new(10, 200).step(10);
let request_timeout_ms_param = IntParam::new(1000, 10000).step(500);
let retry_count_param = IntParam::new(0, 5);
let batch_size_param = IntParam::new(1, 256).log_scale();
let compression_level_param = IntParam::new(0, 9);
let use_http2_param = BoolParam::new();
let load_balancing_param = CategoricalParam::new(vec![
"round_robin",
"least_connections",
"random",
"ip_hash",
]);
// Step 4: Configure optimization
let n_trials = 40;
let concurrency = 4; // Run 4 trials in parallel
@@ -292,7 +284,7 @@ async fn main() -> optimizer::Result<()> {
let start = Instant::now();
// Step 4: Run parallel async optimization
// Step 5: Run parallel async optimization
//
// optimize_parallel_with_sampler:
// - Runs up to `concurrency` trials simultaneously
@@ -303,7 +295,30 @@ async fn main() -> optimizer::Result<()> {
// The "_with_sampler" suffix means the TPE sampler gets access to
// trial history for informed sampling.
study
.optimize_parallel_with_sampler(n_trials, concurrency, objective)
.optimize_parallel_with_sampler(n_trials, concurrency, move |trial| {
let cache_size_mb_param = cache_size_mb_param.clone();
let connection_pool_size_param = connection_pool_size_param.clone();
let request_timeout_ms_param = request_timeout_ms_param.clone();
let retry_count_param = retry_count_param.clone();
let batch_size_param = batch_size_param.clone();
let compression_level_param = compression_level_param.clone();
let use_http2_param = use_http2_param.clone();
let load_balancing_param = load_balancing_param.clone();
async move {
objective(
trial,
&cache_size_mb_param,
&connection_pool_size_param,
&request_timeout_ms_param,
&retry_count_param,
&batch_size_param,
&compression_level_param,
&use_http2_param,
&load_balancing_param,
)
.await
}
})
.await?;
let elapsed = start.elapsed();
examples/ml_hyperparameter_tuning.rs
+71 -70
View File
@@ -23,6 +23,7 @@
use std::ops::ControlFlow;
use optimizer::parameter::{FloatParam, IntParam, Parameter};
use optimizer::sampler::CompletedTrial;
use optimizer::sampler::tpe::TpeSampler;
use optimizer::{Direction, ParamValue, Study, Trial};
@@ -87,34 +88,32 @@ fn evaluate_model(config: &ModelConfig) -> f64 {
/// The objective function that the optimizer calls for each trial.
///
/// This function:
/// 1. Uses `trial.suggest_*()` methods to sample hyperparameter values
/// 2. Builds a model configuration from those values
/// 1. Uses parameter definitions passed as arguments
/// 2. Builds a model configuration from the suggested values
/// 3. Evaluates the model and returns the loss
///
/// The optimizer learns from the results to suggest better parameters
/// in future trials.
fn objective(trial: &mut Trial) -> optimizer::Result<f64> {
// Sample hyperparameters using different strategies:
// Log-scale: Good for parameters spanning multiple orders of magnitude
// The learning rate might be 0.001, 0.01, or 0.1 - log-scale samples evenly across these
let learning_rate = trial.suggest_float_log("learning_rate", 0.001, 0.3)?;
// Regular integer: Uniformly samples from the range [3, 12]
let max_depth = trial.suggest_int("max_depth", 3, 12)?;
// Stepped integer: Only samples multiples of 50 (50, 100, 150, ..., 500)
// Useful when you only want to test specific values
let n_estimators = trial.suggest_int_step("n_estimators", 50, 500, 50)?;
// Regular float: Uniformly samples from [0.5, 1.0]
let subsample = trial.suggest_float("subsample", 0.5, 1.0)?;
let colsample_bytree = trial.suggest_float("colsample_bytree", 0.5, 1.0)?;
// More parameters
let min_child_weight = trial.suggest_int("min_child_weight", 1, 10)?;
let reg_alpha = trial.suggest_float_log("reg_alpha", 1e-3, 10.0)?;
let reg_lambda = trial.suggest_float_log("reg_lambda", 1e-3, 10.0)?;
#[allow(clippy::too_many_arguments)]
fn objective(
trial: &mut Trial,
learning_rate_param: &FloatParam,
max_depth_param: &IntParam,
n_estimators_param: &IntParam,
subsample_param: &FloatParam,
colsample_bytree_param: &FloatParam,
min_child_weight_param: &IntParam,
reg_alpha_param: &FloatParam,
reg_lambda_param: &FloatParam,
) -> optimizer::Result<f64> {
let learning_rate = learning_rate_param.suggest(trial)?;
let max_depth = max_depth_param.suggest(trial)?;
let n_estimators = n_estimators_param.suggest(trial)?;
let subsample = subsample_param.suggest(trial)?;
let colsample_bytree = colsample_bytree_param.suggest(trial)?;
let min_child_weight = min_child_weight_param.suggest(trial)?;
let reg_alpha = reg_alpha_param.suggest(trial)?;
let reg_lambda = reg_lambda_param.suggest(trial)?;
// Build configuration and evaluate
let config = ModelConfig {
@@ -148,35 +147,16 @@ fn objective(trial: &mut Trial) -> optimizer::Result<f64> {
/// Return `ControlFlow::Continue(())` to keep optimizing.
/// Return `ControlFlow::Break(())` to stop early.
fn on_trial_complete(study: &Study<f64>, trial: &CompletedTrial<f64>) -> ControlFlow<()> {
// Helper to extract parameter values
let get_float = |name: &str| -> f64 {
match trial.params.get(name) {
Some(ParamValue::Float(v)) => *v,
_ => 0.0,
// Print trial number and objective value
print!("{:>5} ", study.n_trials());
for value in trial.params.values() {
match value {
ParamValue::Float(v) => print!("{v:>12.5} "),
ParamValue::Int(v) => print!("{v:>12} "),
ParamValue::Categorical(v) => print!("{v:>12} "),
}
};
let get_int = |name: &str| -> i64 {
match trial.params.get(name) {
Some(ParamValue::Int(v)) => *v,
_ => 0,
}
};
// Print progress
println!(
"{:>5} {:>10.5} {:>10} {:>12} {:>10.3} {:>12.3} {:>8} {:>10.4} {:>10.4} {:>12.6}",
study.n_trials(),
get_float("learning_rate"),
get_int("max_depth"),
get_int("n_estimators"),
get_float("subsample"),
get_float("colsample_bytree"),
get_int("min_child_weight"),
get_float("reg_alpha"),
get_float("reg_lambda"),
trial.value,
);
}
println!("{:>12.6}", trial.value);
// Early stopping: if we find an excellent solution, stop early
if trial.value < 0.16 {
@@ -218,21 +198,22 @@ fn main() -> optimizer::Result<()> {
// Print header
println!("Starting hyperparameter optimization...\n");
println!(
"{:>5} {:>10} {:>10} {:>12} {:>10} {:>12} {:>8} {:>10} {:>10} {:>12}",
"Trial",
"LR",
"MaxDepth",
"Estimators",
"Subsample",
"ColSample",
"MinCW",
"Alpha",
"Lambda",
"Loss"
"{:>5} {:>12} (parameters...) {:>12}",
"Trial", "Params", "Loss"
);
println!("{}", "-".repeat(110));
println!("{}", "-".repeat(60));
// Step 3: Run optimization
// Step 3: Define parameter search spaces
let learning_rate_param = FloatParam::new(0.001, 0.3).log_scale();
let max_depth_param = IntParam::new(3, 12);
let n_estimators_param = IntParam::new(50, 500).step(50);
let subsample_param = FloatParam::new(0.5, 1.0);
let colsample_bytree_param = FloatParam::new(0.5, 1.0);
let min_child_weight_param = IntParam::new(1, 10);
let reg_alpha_param = FloatParam::new(1e-3, 10.0).log_scale();
let reg_lambda_param = FloatParam::new(1e-3, 10.0).log_scale();
// Step 4: Run optimization
//
// optimize_with_callback_sampler runs the objective function for up to
// n_trials iterations. After each trial, it calls the callback.
@@ -240,7 +221,23 @@ fn main() -> optimizer::Result<()> {
// history for informed sampling.
let n_trials = 50;
study.optimize_with_callback_sampler(n_trials, objective, on_trial_complete)?;
study.optimize_with_callback_sampler(
n_trials,
|trial| {
objective(
trial,
&learning_rate_param,
&max_depth_param,
&n_estimators_param,
&subsample_param,
&colsample_bytree_param,
&min_child_weight_param,
&reg_alpha_param,
&reg_lambda_param,
)
},
on_trial_complete,
)?;
// Step 4: Get the best result
println!("\n{}", "=".repeat(110));
@@ -252,11 +249,15 @@ fn main() -> optimizer::Result<()> {
println!(" Loss: {:.6}", best.value);
println!(" Parameters:");
for (name, value) in &best.params {
for (id, value) in &best.params {
let label = best
.param_labels
.get(id)
.map_or_else(|| format!("{id}"), |l| l.clone());
match value {
ParamValue::Float(v) => println!(" {name}: {v:.6}"),
ParamValue::Int(v) => println!(" {name}: {v}"),
ParamValue::Categorical(v) => println!(" {name}: category {v}"),
ParamValue::Float(v) => println!(" {label}: {v:.6}"),
ParamValue::Int(v) => println!(" {label}: {v}"),
ParamValue::Categorical(v) => println!(" {label}: category {v}"),
}
}
examples/parameter_api.rs
+55
View File
@@ -0,0 +1,55 @@
use optimizer::parameter::{
BoolParam, CategoricalParam, EnumParam, FloatParam, IntParam, Parameter,
};
use optimizer::{Direction, Study};
use optimizer_derive::Categorical;
#[derive(Clone, Debug, Categorical)]
enum Activation {
Relu,
Sigmoid,
Tanh,
}
fn main() {
let study: Study<f64> = Study::new(Direction::Minimize);
// Define parameters outside the objective function
let lr_param = FloatParam::new(1e-5, 1e-1).log_scale();
let n_layers_param = IntParam::new(1, 5);
let units_param = IntParam::new(32, 512).step(32);
let optimizer_param = CategoricalParam::new(vec!["sgd", "adam", "rmsprop"]);
let activation_param = EnumParam::<Activation>::new();
let batch_size_param = IntParam::new(16, 256).log_scale();
let use_dropout_param = BoolParam::new();
study
.optimize(20, |trial| {
let lr = lr_param.suggest(trial)?;
let n_layers = n_layers_param.suggest(trial)?;
let units = units_param.suggest(trial)?;
let optimizer = optimizer_param.suggest(trial)?;
let use_dropout = use_dropout_param.suggest(trial)?;
let activation = activation_param.suggest(trial)?;
let batch_size = batch_size_param.suggest(trial)?;
// Simulate a loss function
let loss = lr * (n_layers as f64) + (units as f64) * 0.001
- if use_dropout { 0.1 } else { 0.0 };
println!(
"Trial {}: lr={lr:.6}, layers={n_layers}, units={units}, opt={optimizer}, \
dropout={use_dropout}, activation={activation:?}, batch={batch_size} -> loss={loss:.4}",
trial.id()
);
Ok::<_, optimizer::Error>(loss)
})
.unwrap();
let best = study.best_trial().unwrap();
println!("\nBest trial: value={:.4}", best.value);
for (id, label) in &best.param_labels {
println!(" {}: {:?}", label, best.params[id]);
}
}