decisiontree.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include "algo_helper.h"
#include "flatnode.h"
 
#include <string>
#include <vector>
 
namespace ML {
 
namespace DT {
 
struct DecisionTreeParams {
  int max_depth;
  int max_leaves;
  float max_features;
  int max_n_bins;
  int min_samples_leaf;
  int min_samples_split;
  CRITERION split_criterion;
  float min_impurity_decrease = 0.0f;
 
  int max_batch_size;
};
 
void set_tree_params(DecisionTreeParams& params,
                     int cfg_max_depth               = -1,
                     int cfg_max_leaves              = -1,
                     float cfg_max_features          = 1.0f,
                     int cfg_max_n_bins              = 128,
                     int cfg_min_samples_leaf        = 1,
                     int cfg_min_samples_split       = 2,
                     float cfg_min_impurity_decrease = 0.0f,
                     CRITERION cfg_split_criterion   = CRITERION_END,
                     int cfg_max_batch_size          = 4096);
 
template <class T, class L>
struct TreeMetaDataNode {
  int treeid;
  int depth_counter;
  int leaf_counter;
  double train_time;
  std::vector<T> vector_leaf;
  std::vector<SparseTreeNode<T, L>> sparsetree;
  int num_outputs;
};
 
template <class T, class L>
std::string get_tree_summary_text(const TreeMetaDataNode<T, L>* tree);
 
template <class T, class L>
std::string get_tree_text(const TreeMetaDataNode<T, L>* tree);
 
template <class T, class L>
std::string get_tree_json(const TreeMetaDataNode<T, L>* tree);
 
typedef TreeMetaDataNode<float, int> TreeClassifierF;
typedef TreeMetaDataNode<double, int> TreeClassifierD;
typedef TreeMetaDataNode<float, float> TreeRegressorF;
typedef TreeMetaDataNode<double, double> TreeRegressorD;
 
}  // End namespace DT
}  // End namespace ML