Fix default XGBoost n_estimators (#338)

levkk · web-flow · commit 46f493662ca5 · 2022-09-21T14:24:51.000-07:00
diff --git a/pgml-extension/pgml_rust/src/engines/xgboost.rs b/pgml-extension/pgml_rust/src/engines/xgboost.rs
@@ -186,8 +186,8 @@ pub fn xgboost_train(
     let params = parameters::TrainingParametersBuilder::default()
         .dtrain(&dtrain) // dataset to train with
         .boost_rounds(match hyperparams.get("n_estimators") {
-            Some(value) => value.as_u64().unwrap_or(2) as u32,
-            None => 2,
+            Some(value) => value.as_u64().unwrap_or(10) as u32,
+            None => 10,
         }) // number of training iterations
         .booster_params(booster_params) // model parameters
         .evaluation_sets(Some(evaluation_sets)) // optional datasets to evaluate against in each iteration
diff --git a/pgml-extension/pgml_rust/src/lib.rs b/pgml-extension/pgml_rust/src/lib.rs
@@ -3,10 +3,10 @@ extern crate openblas_src;
 extern crate serde;
 
 use once_cell::sync::Lazy; // 1.3.1
+use parking_lot::Mutex;
 use pgx::*;
 use std::collections::HashMap;
 use std::fs;
-use parking_lot::Mutex;
 use xgboost::{Booster, DMatrix};
 
 pub mod api;
diff --git a/pgml-extension/pgml_rust/src/orm/estimator.rs b/pgml-extension/pgml_rust/src/orm/estimator.rs
@@ -1,8 +1,8 @@
+use parking_lot::Mutex;
 use std::collections::HashMap;
 use std::fmt::Debug;
 use std::str::FromStr;
 use std::sync::Arc;
-use parking_lot::Mutex;
 
 use ndarray::{Array1, Array2};
 use once_cell::sync::Lazy;
