fix ResetConfig in FeatureHistogram (#2882)

microsoft · Mar 8, 2020 · 1a48fd2 · 1a48fd2
1 parent 018bdd3
commit 1a48fd2
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/src/treelearner/feature_histogram.hpp b/src/treelearner/feature_histogram.hpp
@@ -59,6 +59,10 @@ class FeatureHistogram {
   void Init(hist_t* data, const FeatureMetainfo* meta) {
     meta_ = meta;
     data_ = data;
+    ResetFunc();
+  }
+
+  void ResetFunc() {
     if (meta_->bin_type == BinType::NumericalBin) {
       FuncForNumrical();
     } else {
@@ -1131,7 +1135,21 @@ class HistogramPool {
   }
 
   void ResetConfig(const Dataset* train_data, const Config* config) {
+    CHECK_GT(train_data->num_features(), 0);
+    const Config* old_config = feature_metas_[0].config;
     SetFeatureInfo<false, true>(train_data, config, &feature_metas_);
+    // if need to reset the function pointers
+    if (old_config->lambda_l1 != config->lambda_l1 ||
+        old_config->monotone_constraints != config->monotone_constraints ||
+        old_config->extra_trees != config->extra_trees ||
+        old_config->max_delta_step != config->max_delta_step) {
+#pragma omp parallel for schedule(static)
+      for (int i = 0; i < cache_size_; ++i) {
+        for (int j = 0; j < train_data->num_features(); ++j) {
+          pool_[i][j].ResetFunc();
+        }
+      }
+    }
   }
 
   /*!