templatize refresher

This commit is contained in:
tqchen 2014-08-24 15:22:11 -07:00
parent f71b732e7a
commit ba9fbd380c
2 changed files with 9 additions and 9 deletions

View File

@ -60,7 +60,7 @@ namespace tree {
template<typename FMatrix>
inline IUpdater<FMatrix>* CreateUpdater(const char *name) {
if (!strcmp(name, "prune")) return new TreePruner<FMatrix>();
if (!strcmp(name, "refresh")) return new TreeRefresher<FMatrix>();
if (!strcmp(name, "refresh")) return new TreeRefresher<FMatrix, GradStats>();
if (!strcmp(name, "grow_colmaker")) return new ColMaker<FMatrix, GradStats>();
utils::Error("unknown updater:%s", name);
return NULL;

View File

@ -13,7 +13,7 @@
namespace xgboost {
namespace tree {
/*! \brief pruner that prunes a tree after growing finishs */
template<typename FMatrix>
template<typename FMatrix, typename TStats>
class TreeRefresher: public IUpdater<FMatrix> {
public:
virtual ~TreeRefresher(void) {}
@ -30,7 +30,7 @@ class TreeRefresher: public IUpdater<FMatrix> {
// number of threads
int nthread;
// thread temporal space
std::vector< std::vector<GradStats> > stemp;
std::vector< std::vector<TStats> > stemp;
std::vector<RegTree::FVec> fvec_temp;
// setup temp space for each thread
#pragma omp parallel
@ -38,14 +38,14 @@ class TreeRefresher: public IUpdater<FMatrix> {
nthread = omp_get_num_threads();
}
fvec_temp.resize(nthread, RegTree::FVec());
stemp.resize(trees.size() * nthread, std::vector<GradStats>());
stemp.resize(trees.size() * nthread, std::vector<TStats>());
#pragma omp parallel
{
int tid = omp_get_thread_num();
for (size_t i = 0; i < trees.size(); ++i) {
std::vector<GradStats> &vec = stemp[tid * trees.size() + i];
std::vector<TStats> &vec = stemp[tid * trees.size() + i];
vec.resize(trees[i]->param.num_nodes);
std::fill(vec.begin(), vec.end(), GradStats());
std::fill(vec.begin(), vec.end(), TStats());
}
fvec_temp[tid].Init(trees[0]->param.num_feature);
}
@ -97,8 +97,8 @@ class TreeRefresher: public IUpdater<FMatrix> {
const std::vector<bst_gpair> &gpair,
const BoosterInfo &info,
const bst_uint ridx,
std::vector<GradStats> *p_gstats) {
std::vector<GradStats> &gstats = *p_gstats;
std::vector<TStats> *p_gstats) {
std::vector<TStats> &gstats = *p_gstats;
// start from groups that belongs to current data
int pid = static_cast<int>(info.GetRoot(ridx));
gstats[pid].Add(gpair, info, ridx);
@ -109,7 +109,7 @@ class TreeRefresher: public IUpdater<FMatrix> {
gstats[pid].Add(gpair, info, ridx);
}
}
inline void Refresh(const std::vector<GradStats> &gstats,
inline void Refresh(const std::vector<TStats> &gstats,
int nid, RegTree *p_tree) {
RegTree &tree = *p_tree;
tree.stat(nid).base_weight = gstats[nid].CalcWeight(param);