perf(tpe): eliminate unnecessary Vec allocations in KDE/TPE sampling
- In-place log transformation in sample_tpe_float (saves 2 allocs per log-scale float sample) - Pass owned Vec<i64> to sample_tpe_int to reduce peak memory during int-to-float conversion - Stack-allocate categorical count arrays for <=32 choices in sample_tpe_categorical
This commit is contained in:
+44
-11
@@ -263,8 +263,20 @@ impl MotpeSampler {
|
|||||||
let (internal_low, internal_high, good_internal, bad_internal) = if log_scale {
|
let (internal_low, internal_high, good_internal, bad_internal) = if log_scale {
|
||||||
let i_low = low.ln();
|
let i_low = low.ln();
|
||||||
let i_high = high.ln();
|
let i_high = high.ln();
|
||||||
let g: Vec<f64> = good_values.iter().map(|&v| v.ln()).collect();
|
let g = {
|
||||||
let b: Vec<f64> = bad_values.iter().map(|&v| v.ln()).collect();
|
let mut v = good_values;
|
||||||
|
for x in &mut v {
|
||||||
|
*x = x.ln();
|
||||||
|
}
|
||||||
|
v
|
||||||
|
};
|
||||||
|
let b = {
|
||||||
|
let mut v = bad_values;
|
||||||
|
for x in &mut v {
|
||||||
|
*x = x.ln();
|
||||||
|
}
|
||||||
|
v
|
||||||
|
};
|
||||||
(i_low, i_high, g, b)
|
(i_low, i_high, g, b)
|
||||||
} else {
|
} else {
|
||||||
(low, high, good_values, bad_values)
|
(low, high, good_values, bad_values)
|
||||||
@@ -339,12 +351,12 @@ impl MotpeSampler {
|
|||||||
high: i64,
|
high: i64,
|
||||||
log_scale: bool,
|
log_scale: bool,
|
||||||
step: Option<i64>,
|
step: Option<i64>,
|
||||||
good_values: &[i64],
|
good_values: Vec<i64>,
|
||||||
bad_values: &[i64],
|
bad_values: Vec<i64>,
|
||||||
rng: &mut fastrand::Rng,
|
rng: &mut fastrand::Rng,
|
||||||
) -> i64 {
|
) -> i64 {
|
||||||
let good_floats: Vec<f64> = good_values.iter().map(|&v| v as f64).collect();
|
let good_floats: Vec<f64> = good_values.into_iter().map(|v| v as f64).collect();
|
||||||
let bad_floats: Vec<f64> = bad_values.iter().map(|&v| v as f64).collect();
|
let bad_floats: Vec<f64> = bad_values.into_iter().map(|v| v as f64).collect();
|
||||||
|
|
||||||
let float_value = self.sample_tpe_float(
|
let float_value = self.sample_tpe_float(
|
||||||
low as f64,
|
low as f64,
|
||||||
@@ -375,9 +387,30 @@ impl MotpeSampler {
|
|||||||
bad_indices: &[usize],
|
bad_indices: &[usize],
|
||||||
rng: &mut fastrand::Rng,
|
rng: &mut fastrand::Rng,
|
||||||
) -> usize {
|
) -> usize {
|
||||||
let mut good_counts = vec![0usize; n_choices];
|
// Stack-allocate for the common case (<=32 choices), heap for rare large cases
|
||||||
let mut bad_counts = vec![0usize; n_choices];
|
let mut good_buf = [0usize; 32];
|
||||||
|
let mut bad_buf = [0usize; 32];
|
||||||
|
let mut weight_buf = [0.0f64; 32];
|
||||||
|
|
||||||
|
let mut good_vec;
|
||||||
|
let mut bad_vec;
|
||||||
|
let mut weight_vec;
|
||||||
|
|
||||||
|
let (good_counts, bad_counts, weights): (&mut [usize], &mut [usize], &mut [f64]) =
|
||||||
|
if n_choices <= 32 {
|
||||||
|
(
|
||||||
|
&mut good_buf[..n_choices],
|
||||||
|
&mut bad_buf[..n_choices],
|
||||||
|
&mut weight_buf[..n_choices],
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
good_vec = vec![0usize; n_choices];
|
||||||
|
bad_vec = vec![0usize; n_choices];
|
||||||
|
weight_vec = vec![0.0f64; n_choices];
|
||||||
|
(&mut good_vec, &mut bad_vec, &mut weight_vec)
|
||||||
|
};
|
||||||
|
|
||||||
|
// Count occurrences in good and bad groups
|
||||||
for &idx in good_indices {
|
for &idx in good_indices {
|
||||||
if idx < n_choices {
|
if idx < n_choices {
|
||||||
good_counts[idx] += 1;
|
good_counts[idx] += 1;
|
||||||
@@ -393,7 +426,7 @@ impl MotpeSampler {
|
|||||||
let good_total = good_indices.len() as f64 + n_choices as f64;
|
let good_total = good_indices.len() as f64 + n_choices as f64;
|
||||||
let bad_total = bad_indices.len() as f64 + n_choices as f64;
|
let bad_total = bad_indices.len() as f64 + n_choices as f64;
|
||||||
|
|
||||||
let mut weights = vec![0.0f64; n_choices];
|
// Calculate l(x)/g(x) ratio for each category
|
||||||
for i in 0..n_choices {
|
for i in 0..n_choices {
|
||||||
let l_prob = (good_counts[i] as f64 + 1.0) / good_total;
|
let l_prob = (good_counts[i] as f64 + 1.0) / good_total;
|
||||||
let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total;
|
let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total;
|
||||||
@@ -521,8 +554,8 @@ impl MultiObjectiveSampler for MotpeSampler {
|
|||||||
d.high,
|
d.high,
|
||||||
d.log_scale,
|
d.log_scale,
|
||||||
d.step,
|
d.step,
|
||||||
&good_values,
|
good_values,
|
||||||
&bad_values,
|
bad_values,
|
||||||
&mut rng,
|
&mut rng,
|
||||||
);
|
);
|
||||||
ParamValue::Int(value)
|
ParamValue::Int(value)
|
||||||
|
|||||||
@@ -1356,8 +1356,8 @@ impl MultivariateTpeSampler {
|
|||||||
d.high,
|
d.high,
|
||||||
d.log_scale,
|
d.log_scale,
|
||||||
d.step,
|
d.step,
|
||||||
&good_values,
|
good_values,
|
||||||
&bad_values,
|
bad_values,
|
||||||
rng,
|
rng,
|
||||||
);
|
);
|
||||||
ParamValue::Int(value)
|
ParamValue::Int(value)
|
||||||
@@ -1412,8 +1412,20 @@ impl MultivariateTpeSampler {
|
|||||||
let (internal_low, internal_high, good_internal, bad_internal) = if log_scale {
|
let (internal_low, internal_high, good_internal, bad_internal) = if log_scale {
|
||||||
let i_low = low.ln();
|
let i_low = low.ln();
|
||||||
let i_high = high.ln();
|
let i_high = high.ln();
|
||||||
let g: Vec<f64> = good_values.iter().map(|&v| v.ln()).collect();
|
let g = {
|
||||||
let b: Vec<f64> = bad_values.iter().map(|&v| v.ln()).collect();
|
let mut v = good_values;
|
||||||
|
for x in &mut v {
|
||||||
|
*x = x.ln();
|
||||||
|
}
|
||||||
|
v
|
||||||
|
};
|
||||||
|
let b = {
|
||||||
|
let mut v = bad_values;
|
||||||
|
for x in &mut v {
|
||||||
|
*x = x.ln();
|
||||||
|
}
|
||||||
|
v
|
||||||
|
};
|
||||||
(i_low, i_high, g, b)
|
(i_low, i_high, g, b)
|
||||||
} else {
|
} else {
|
||||||
(low, high, good_values, bad_values)
|
(low, high, good_values, bad_values)
|
||||||
@@ -1487,13 +1499,13 @@ impl MultivariateTpeSampler {
|
|||||||
high: i64,
|
high: i64,
|
||||||
log_scale: bool,
|
log_scale: bool,
|
||||||
step: Option<i64>,
|
step: Option<i64>,
|
||||||
good_values: &[i64],
|
good_values: Vec<i64>,
|
||||||
bad_values: &[i64],
|
bad_values: Vec<i64>,
|
||||||
rng: &mut fastrand::Rng,
|
rng: &mut fastrand::Rng,
|
||||||
) -> i64 {
|
) -> i64 {
|
||||||
// Convert to floats for KDE
|
// Convert to floats for KDE
|
||||||
let good_floats: Vec<f64> = good_values.iter().map(|&v| v as f64).collect();
|
let good_floats: Vec<f64> = good_values.into_iter().map(|v| v as f64).collect();
|
||||||
let bad_floats: Vec<f64> = bad_values.iter().map(|&v| v as f64).collect();
|
let bad_floats: Vec<f64> = bad_values.into_iter().map(|v| v as f64).collect();
|
||||||
|
|
||||||
// Use float TPE sampling
|
// Use float TPE sampling
|
||||||
let float_value = self.sample_tpe_float(
|
let float_value = self.sample_tpe_float(
|
||||||
@@ -1598,10 +1610,30 @@ impl MultivariateTpeSampler {
|
|||||||
bad_indices: &[usize],
|
bad_indices: &[usize],
|
||||||
rng: &mut fastrand::Rng,
|
rng: &mut fastrand::Rng,
|
||||||
) -> usize {
|
) -> usize {
|
||||||
// Count occurrences in good and bad groups
|
// Stack-allocate for the common case (<=32 choices), heap for rare large cases
|
||||||
let mut good_counts = vec![0usize; n_choices];
|
let mut good_buf = [0usize; 32];
|
||||||
let mut bad_counts = vec![0usize; n_choices];
|
let mut bad_buf = [0usize; 32];
|
||||||
|
let mut weight_buf = [0.0f64; 32];
|
||||||
|
|
||||||
|
let mut good_vec;
|
||||||
|
let mut bad_vec;
|
||||||
|
let mut weight_vec;
|
||||||
|
|
||||||
|
let (good_counts, bad_counts, weights): (&mut [usize], &mut [usize], &mut [f64]) =
|
||||||
|
if n_choices <= 32 {
|
||||||
|
(
|
||||||
|
&mut good_buf[..n_choices],
|
||||||
|
&mut bad_buf[..n_choices],
|
||||||
|
&mut weight_buf[..n_choices],
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
good_vec = vec![0usize; n_choices];
|
||||||
|
bad_vec = vec![0usize; n_choices];
|
||||||
|
weight_vec = vec![0.0f64; n_choices];
|
||||||
|
(&mut good_vec, &mut bad_vec, &mut weight_vec)
|
||||||
|
};
|
||||||
|
|
||||||
|
// Count occurrences in good and bad groups
|
||||||
for &idx in good_indices {
|
for &idx in good_indices {
|
||||||
if idx < n_choices {
|
if idx < n_choices {
|
||||||
good_counts[idx] += 1;
|
good_counts[idx] += 1;
|
||||||
@@ -1618,7 +1650,6 @@ impl MultivariateTpeSampler {
|
|||||||
let bad_total = bad_indices.len() as f64 + n_choices as f64;
|
let bad_total = bad_indices.len() as f64 + n_choices as f64;
|
||||||
|
|
||||||
// Calculate l(x)/g(x) ratio for each category
|
// Calculate l(x)/g(x) ratio for each category
|
||||||
let mut weights = vec![0.0f64; n_choices];
|
|
||||||
for i in 0..n_choices {
|
for i in 0..n_choices {
|
||||||
let l_prob = (good_counts[i] as f64 + 1.0) / good_total;
|
let l_prob = (good_counts[i] as f64 + 1.0) / good_total;
|
||||||
let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total;
|
let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total;
|
||||||
|
|||||||
+43
-12
@@ -373,8 +373,20 @@ impl TpeSampler {
|
|||||||
let (internal_low, internal_high, good_internal, bad_internal) = if log_scale {
|
let (internal_low, internal_high, good_internal, bad_internal) = if log_scale {
|
||||||
let i_low = low.ln();
|
let i_low = low.ln();
|
||||||
let i_high = high.ln();
|
let i_high = high.ln();
|
||||||
let g: Vec<f64> = good_values.iter().map(|&v| v.ln()).collect();
|
let g = {
|
||||||
let b: Vec<f64> = bad_values.iter().map(|&v| v.ln()).collect();
|
let mut v = good_values;
|
||||||
|
for x in &mut v {
|
||||||
|
*x = x.ln();
|
||||||
|
}
|
||||||
|
v
|
||||||
|
};
|
||||||
|
let b = {
|
||||||
|
let mut v = bad_values;
|
||||||
|
for x in &mut v {
|
||||||
|
*x = x.ln();
|
||||||
|
}
|
||||||
|
v
|
||||||
|
};
|
||||||
(i_low, i_high, g, b)
|
(i_low, i_high, g, b)
|
||||||
} else {
|
} else {
|
||||||
(low, high, good_values, bad_values)
|
(low, high, good_values, bad_values)
|
||||||
@@ -454,13 +466,13 @@ impl TpeSampler {
|
|||||||
high: i64,
|
high: i64,
|
||||||
log_scale: bool,
|
log_scale: bool,
|
||||||
step: Option<i64>,
|
step: Option<i64>,
|
||||||
good_values: &[i64],
|
good_values: Vec<i64>,
|
||||||
bad_values: &[i64],
|
bad_values: Vec<i64>,
|
||||||
rng: &mut fastrand::Rng,
|
rng: &mut fastrand::Rng,
|
||||||
) -> i64 {
|
) -> i64 {
|
||||||
// Convert to floats for KDE
|
// Convert to floats for KDE
|
||||||
let good_floats: Vec<f64> = good_values.iter().map(|&v| v as f64).collect();
|
let good_floats: Vec<f64> = good_values.into_iter().map(|v| v as f64).collect();
|
||||||
let bad_floats: Vec<f64> = bad_values.iter().map(|&v| v as f64).collect();
|
let bad_floats: Vec<f64> = bad_values.into_iter().map(|v| v as f64).collect();
|
||||||
|
|
||||||
// Use float TPE sampling
|
// Use float TPE sampling
|
||||||
let float_value = self.sample_tpe_float(
|
let float_value = self.sample_tpe_float(
|
||||||
@@ -497,10 +509,30 @@ impl TpeSampler {
|
|||||||
bad_indices: &[usize],
|
bad_indices: &[usize],
|
||||||
rng: &mut fastrand::Rng,
|
rng: &mut fastrand::Rng,
|
||||||
) -> usize {
|
) -> usize {
|
||||||
// Count occurrences in good and bad groups
|
// Stack-allocate for the common case (<=32 choices), heap for rare large cases
|
||||||
let mut good_counts = vec![0usize; n_choices];
|
let mut good_buf = [0usize; 32];
|
||||||
let mut bad_counts = vec![0usize; n_choices];
|
let mut bad_buf = [0usize; 32];
|
||||||
|
let mut weight_buf = [0.0f64; 32];
|
||||||
|
|
||||||
|
let mut good_vec;
|
||||||
|
let mut bad_vec;
|
||||||
|
let mut weight_vec;
|
||||||
|
|
||||||
|
let (good_counts, bad_counts, weights): (&mut [usize], &mut [usize], &mut [f64]) =
|
||||||
|
if n_choices <= 32 {
|
||||||
|
(
|
||||||
|
&mut good_buf[..n_choices],
|
||||||
|
&mut bad_buf[..n_choices],
|
||||||
|
&mut weight_buf[..n_choices],
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
good_vec = vec![0usize; n_choices];
|
||||||
|
bad_vec = vec![0usize; n_choices];
|
||||||
|
weight_vec = vec![0.0f64; n_choices];
|
||||||
|
(&mut good_vec, &mut bad_vec, &mut weight_vec)
|
||||||
|
};
|
||||||
|
|
||||||
|
// Count occurrences in good and bad groups
|
||||||
for &idx in good_indices {
|
for &idx in good_indices {
|
||||||
if idx < n_choices {
|
if idx < n_choices {
|
||||||
good_counts[idx] += 1;
|
good_counts[idx] += 1;
|
||||||
@@ -517,7 +549,6 @@ impl TpeSampler {
|
|||||||
let bad_total = bad_indices.len() as f64 + n_choices as f64;
|
let bad_total = bad_indices.len() as f64 + n_choices as f64;
|
||||||
|
|
||||||
// Calculate l(x)/g(x) ratio for each category
|
// Calculate l(x)/g(x) ratio for each category
|
||||||
let mut weights = vec![0.0f64; n_choices];
|
|
||||||
for i in 0..n_choices {
|
for i in 0..n_choices {
|
||||||
let l_prob = (good_counts[i] as f64 + 1.0) / good_total;
|
let l_prob = (good_counts[i] as f64 + 1.0) / good_total;
|
||||||
let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total;
|
let g_prob = (bad_counts[i] as f64 + 1.0) / bad_total;
|
||||||
@@ -967,8 +998,8 @@ impl Sampler for TpeSampler {
|
|||||||
d.high,
|
d.high,
|
||||||
d.log_scale,
|
d.log_scale,
|
||||||
d.step,
|
d.step,
|
||||||
&good_values,
|
good_values,
|
||||||
&bad_values,
|
bad_values,
|
||||||
&mut rng,
|
&mut rng,
|
||||||
);
|
);
|
||||||
ParamValue::Int(value)
|
ParamValue::Int(value)
|
||||||
|
|||||||
Reference in New Issue
Block a user