decisiontree.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2019-2023, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#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