silvic_core/
boosting.rs

1//! The boosting loop (LightGBM 3.3.5 `GBDT::Train`/`GBDT::TrainOneIter`, `src/boosting/gbdt.cpp`).
2//!
3//! Each iteration computes gradients, reduces their maxima over ranks and quantizes them (quantized mode), draws
4//! the bag, grows a tree, renews its leaves from exact sums (quantized mode), shrinks it, adds it to every dataset's
5//! scores and evaluates the validation sets. As in LightGBM the initial score (`BoostFromAverage`) is added to the
6//! scores up front and folded into the first tree as a bias; a first tree that cannot split becomes a constant tree
7//! holding it, and any later tree that cannot split ends training without being added.
8//!
9//! [`train_from`] continues an existing model (LightGBM `init_model`): the scores start from the initial score plus
10//! the model's trees replayed on the binned rows (`replay`), the first tree's bias taken back off, so they are the
11//! f32 scores the original training had; the iteration-keyed draws (bagging round, feature sample, quantization)
12//! continue at the model's tree count. A model checkpointed by silvic and continued on the same rows and parameters
13//! therefore trains the same trees as one uninterrupted run.
14
15use std::sync::Arc;
16
17use rayon::ThreadPool;
18
19use crate::backend::{Backend, BagSpec, ObjectiveSetup};
20use crate::collective::Collective;
21use crate::config::Params;
22use crate::data::BinnedDataset;
23use crate::error::{Result, SilvicError};
24use crate::learner::{GrownTree, Learner, TreeSpec};
25use crate::metric::{self, EvalResult};
26use crate::model::{Model, ModelObjective};
27use crate::objective;
28use crate::quantize::QuantSpec;
29use crate::replay;
30use crate::split::{self, HistScale, K_EPSILON};
31use crate::tree::Tree;
32
33/// What the iteration callback sees after each tree.
34#[derive(Debug, Clone, PartialEq)]
35pub struct IterationReport<'a> {
36    /// Zero-based and global: an `init_model`'s trees count, so the first new tree after `k` of them is `k`.
37    pub iteration: usize,
38    pub evals: &'a [EvalResult],
39    head: &'a Model,
40    init_trees: &'a [Tree],
41    trees: &'a [Tree],
42}
43
44impl IterationReport<'_> {
45    /// The model as of this iteration: the `init_model`'s trees, then every tree trained so far. Built on demand
46    /// (O(trees)), for checkpoints.
47    pub fn model(&self) -> Model {
48        let mut trees = Vec::with_capacity(self.init_trees.len() + self.trees.len());
49        trees.extend_from_slice(self.init_trees);
50        trees.extend_from_slice(self.trees);
51        Model {
52            trees,
53            ..self.head.clone()
54        }
55    }
56}
57
58/// Trains `params.num_iterations` trees on `train`, evaluating `valid` after each. Every rank calls this with its own
59/// rows; every rank returns the same model. An error from `on_iteration` stops training and is returned.
60///
61/// With `num_threads > 0`, split finding, metrics and the backend's host work run on a pool of that many threads
62/// (handed to the backend through `Backend::set_thread_pool`); otherwise on rayon's global pool.
63pub fn train(
64    params: &Params,
65    train: &BinnedDataset,
66    valid: &[(&str, &BinnedDataset)],
67    backend: &mut dyn Backend,
68    collective: &mut dyn Collective,
69    on_iteration: &mut dyn FnMut(&IterationReport<'_>) -> Result<()>,
70) -> Result<Model> {
71    train_from(params, None, train, valid, backend, collective, on_iteration)
72}
73
74/// [`train`], continuing `init_model` when there is one (LightGBM `init_model`): `params.num_iterations` new trees
75/// follow its trees, and the returned model holds both. Its `num_iterations` parameter is the total, so a run
76/// resumed from a checkpoint writes the same model text as the uninterrupted run; `pandas_categorical` is kept from
77/// `init_model`.
78///
79/// `init_model` must have the objective (and sigmoid) and the feature count of `params` and `train`, else
80/// [`SilvicError::Config`]. A model trained by silvic on bins from the same sample continues exactly; a model
81/// trained elsewhere (other bins) continues from approximately replayed scores (see `replay`). The initial score is
82/// recomputed from `params` and `train` and taken off the first tree's leaves, as `boost_from_average` folded it in.
83pub fn train_from(
84    params: &Params,
85    init_model: Option<&Model>,
86    train: &BinnedDataset,
87    valid: &[(&str, &BinnedDataset)],
88    backend: &mut dyn Backend,
89    collective: &mut dyn Collective,
90    on_iteration: &mut dyn FnMut(&IterationReport<'_>) -> Result<()>,
91) -> Result<Model> {
92    let pool = thread_pool(params.num_threads)?;
93    if let Some(pool) = &pool {
94        backend.set_thread_pool(Arc::clone(pool));
95    }
96    let setup = objective::setup(params, train)?;
97    let init_trees: &[Tree] = init_model.map_or(&[], |m| &m.trees);
98    let head = model_head(params, &setup, train, init_model)?;
99    let valid_sets: Vec<&BinnedDataset> = valid.iter().map(|(_, data)| *data).collect();
100    backend.load(train, &valid_sets, &setup)?;
101    let init_score = install(pool.as_deref(), || {
102        objective::init_score(params, train, &mut *collective)
103    })?;
104    // `GBDT::BoostFromAverage` ignores initial scores within kEpsilon of 0.
105    let init_score = if init_score.abs() > K_EPSILON { init_score } else { 0.0 };
106    for dataset in 0..=valid.len() {
107        backend.init_scores(dataset, init_score)?;
108    }
109    for (index, tree) in init_trees.iter().enumerate() {
110        let bias = if index == 0 { init_score } else { 0.0 };
111        backend.add_tree(&replay::device_tree(tree, train, bias)?)?;
112    }
113    tracing::debug!(backend = %backend.name(), init_score, init_trees = init_trees.len(), "training started");
114    let mut trainer = Trainer {
115        params,
116        train,
117        valid,
118        learner: Learner::new(params, &train.features, train.total_bins, pool.as_deref()),
119        pool: pool.as_deref(),
120        constant_hessian: objective::constant_hessian(&setup, train.weight.as_deref()),
121        init_score,
122        prior_trees: init_trees.len(),
123        trees: Vec::new(),
124    };
125    let start = init_trees.len();
126    for iteration in start..start + params.num_iterations {
127        if !trainer.train_one_iter(backend, collective, iteration)? {
128            break;
129        }
130        let evals = trainer.evaluate(backend, collective)?;
131        tracing::debug!(
132            iteration,
133            leaves = trainer.trees.last().map_or(0, |t| t.num_leaves),
134            ?evals,
135            "tree added"
136        );
137        on_iteration(&IterationReport {
138            iteration,
139            evals: &evals,
140            head: &head,
141            init_trees,
142            trees: &trainer.trees,
143        })?;
144    }
145    let mut model = head;
146    model.trees = init_trees.to_vec();
147    model.trees.append(&mut trainer.trees);
148    Ok(model)
149}
150
151/// The model without trees: objective, features and parameters, checked against `init_model`.
152fn model_head(
153    params: &Params,
154    setup: &ObjectiveSetup,
155    train: &BinnedDataset,
156    init_model: Option<&Model>,
157) -> Result<Model> {
158    let objective = match setup {
159        ObjectiveSetup::Regression => ModelObjective::Regression,
160        ObjectiveSetup::Binary { sigmoid } => ModelObjective::Binary { sigmoid: *sigmoid },
161        ObjectiveSetup::LambdaRank { .. } => ModelObjective::LambdaRank,
162    };
163    let mut parameters = params.model_parameters();
164    if let Some(init) = init_model {
165        if init.objective != objective {
166            return Err(SilvicError::config(format!(
167                "init_model has objective {:?}, the parameters train {objective:?}",
168                init.objective
169            )));
170        }
171        if init.num_features() != train.features.len() {
172            return Err(SilvicError::config(format!(
173                "init_model has {} features, the training data has {}",
174                init.num_features(),
175                train.features.len()
176            )));
177        }
178        let total = init.num_trees() + params.num_iterations;
179        if let Some((_, value)) = parameters.iter_mut().find(|(key, _)| key == "num_iterations") {
180            *value = total.to_string();
181        }
182    }
183    Ok(Model {
184        objective,
185        feature_names: train.features.iter().map(|f| f.name.clone()).collect(),
186        feature_infos: train.features.iter().map(|f| f.bins.info_string()).collect(),
187        trees: Vec::new(),
188        parameters,
189        pandas_categorical: init_model.and_then(|m| m.pandas_categorical.clone()),
190    })
191}
192
193fn thread_pool(num_threads: usize) -> Result<Option<Arc<ThreadPool>>> {
194    if num_threads == 0 {
195        return Ok(None);
196    }
197    rayon::ThreadPoolBuilder::new()
198        .num_threads(num_threads)
199        .thread_name(|index| format!("silvic-{index}"))
200        .build()
201        .map(|pool| Some(Arc::new(pool)))
202        .map_err(|e| SilvicError::config(format!("cannot start {num_threads} threads: {e}")))
203}
204
205fn install<R: Send>(pool: Option<&ThreadPool>, op: impl FnOnce() -> R + Send) -> R {
206    match pool {
207        Some(pool) => pool.install(op),
208        None => op(),
209    }
210}
211
212struct Trainer<'a> {
213    params: &'a Params,
214    train: &'a BinnedDataset,
215    valid: &'a [(&'a str, &'a BinnedDataset)],
216    learner: Learner<'a>,
217    pool: Option<&'a ThreadPool>,
218    constant_hessian: Option<f64>,
219    init_score: f64,
220    /// Trees of the `init_model`, before `trees`.
221    prior_trees: usize,
222    trees: Vec<Tree>,
223}
224
225impl Trainer<'_> {
226    /// `GBDT::TrainOneIter`; returns whether training continues.
227    fn train_one_iter(
228        &mut self,
229        backend: &mut dyn Backend,
230        collective: &mut dyn Collective,
231        iteration: usize,
232    ) -> Result<bool> {
233        let params = self.params;
234        let stats = backend.compute_gradients()?;
235        let mut maxima = [stats.max_abs_grad, stats.max_hess];
236        let scale = if params.use_quantized_grad {
237            collective.allreduce_max_f64(&mut maxima)?;
238            let spec = QuantSpec::new(
239                maxima[0],
240                maxima[1],
241                self.constant_hessian,
242                params.num_grad_quant_bins,
243                params.stochastic_rounding,
244                params.data_random_seed,
245                iteration as u32,
246            );
247            backend.quantize(&spec)?;
248            HistScale {
249                grad: spec.grad_scale,
250                hess: spec.hess_scale,
251            }
252        } else {
253            HistScale::UNIT
254        };
255        let local_rows = backend.reset_partition(&self.bag(iteration))?;
256        let mut root_count = [i64::from(local_rows)];
257        collective.allreduce_sum_i64(&mut root_count)?;
258        let feature_mask = split::feature_mask(
259            &self.train.features,
260            params.feature_fraction,
261            params.feature_fraction_seed,
262            iteration as u32,
263        );
264        let spec = TreeSpec {
265            quantized: params.use_quantized_grad,
266            scale,
267            feature_mask: &feature_mask,
268            root: crate::backend::LeafRange {
269                begin: 0,
270                count: local_rows,
271            },
272            root_count: root_count[0],
273        };
274        let mut grown = self.learner.grow(backend, collective, spec)?;
275        if grown.tree.num_leaves == 1 {
276            self.single_leaf(iteration);
277            return Ok(false);
278        }
279        if params.use_quantized_grad && params.quant_train_renew_leaf {
280            self.learner
281                .renew_leaves(backend, collective, &mut grown, maxima[0], maxima[1])?;
282        }
283        self.add(backend, grown)?;
284        Ok(true)
285    }
286
287    /// LightGBM bags every `bagging_freq` iterations; the bag of a round is a pure function of the round.
288    fn bag(&self, iteration: usize) -> BagSpec {
289        let params = self.params;
290        if params.bagging_fraction < 1.0 && params.bagging_freq > 0 {
291            BagSpec::Bernoulli {
292                seed: params.bagging_seed,
293                round: (iteration / params.bagging_freq) as u32,
294                fraction: params.bagging_fraction,
295            }
296        } else {
297            BagSpec::All
298        }
299    }
300
301    /// A tree that cannot split: the first one becomes a constant tree with the initial score, as in
302    /// `GBDT::TrainOneIter`; a later one is dropped. Either way training ends.
303    fn single_leaf(&mut self, iteration: usize) {
304        tracing::debug!(iteration, "no split with positive gain; training stops");
305        if self.prior_trees + self.trees.len() == 0 {
306            let mut tree = Tree::new();
307            tree.as_constant(self.init_score);
308            self.trees.push(tree);
309        }
310    }
311
312    /// Shrinkage, the first tree's bias, and the score update (without the bias, already in the scores).
313    fn add(&mut self, backend: &mut dyn Backend, mut grown: GrownTree) -> Result<()> {
314        grown.tree.shrink(self.params.learning_rate);
315        let mut device = grown.device_tree();
316        if self.prior_trees + self.trees.len() == 0 && self.init_score != 0.0 {
317            grown.tree.add_bias(self.init_score);
318            // As `train_from` replays it, so a resumed run adds bit for bit the same scores.
319            device.leaf_values = replay::leaf_values(&grown.tree, self.init_score);
320        }
321        backend.add_tree(&device)?;
322        self.trees.push(grown.tree);
323        Ok(())
324    }
325
326    fn evaluate(&self, backend: &mut dyn Backend, collective: &mut dyn Collective) -> Result<Vec<EvalResult>> {
327        let mut evals = Vec::new();
328        for (index, (name, data)) in self.valid.iter().enumerate() {
329            let scores = backend.scores(index + 1)?;
330            let collective = &mut *collective;
331            evals.extend(install(self.pool, || {
332                metric::evaluate(self.params, name, &scores, data, collective)
333            })?);
334        }
335        Ok(evals)
336    }
337}
silvic_core/boosting.rs

silvic_core/
boosting.rs
