tree_learner.cpp 666 Bytes
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
3
4
5
6
7
#include <LightGBM/tree_learner.h>

#include "serial_tree_learner.h"
#include "parallel_tree_learner.h"

namespace LightGBM {

8
9
TreeLearner* TreeLearner::CreateTreeLearner(const std::string& type, const TreeConfig* tree_config) {
  if (type == std::string("serial")) {
Guolin Ke's avatar
Guolin Ke committed
10
    return new SerialTreeLearner(tree_config);
11
  } else if (type == std::string("feature")) {
Guolin Ke's avatar
Guolin Ke committed
12
    return new FeatureParallelTreeLearner(tree_config);
13
  } else if (type == std::string("data")) {
Guolin Ke's avatar
Guolin Ke committed
14
    return new DataParallelTreeLearner(tree_config);
15
  } else if (type == std::string("voting")) {
Guolin Ke's avatar
Guolin Ke committed
16
    return new VotingParallelTreeLearner(tree_config);
Guolin Ke's avatar
Guolin Ke committed
17
18
19
20
21
  }
  return nullptr;
}

}  // namespace LightGBM