XGBoost源碼分析之單機多線程的實現

上篇文章主要通過論文閱讀、數學推導,基本掌握了XGBoost的原理。於是開始閱讀XGBoost源碼,並總結了幾處自己認爲比較重要的方面。如有錯誤,請指正:

 

1. 總體框架:


cli_main.cc 是程序的入口,main函數所在的文件。除了有main函數以外,還有訓練參數的結構體。如模型保存路徑,數據路徑,迭代次數等。這個源碼註釋的很清楚,不再贅述。

通過調用main() -> CLIRunTask() -> CLITrain(),這裏我們主要看函數CLITrain()的流程

 

CLITrain主要分了幾步:

1.加載數據

2.初始化learner

3.調用learner->InitModel()、learner->Configure初始化Model,確定目標函數和模型,初始化目標函數結構體和模型

4.根據模型參數param.num_round迭代調用UpdateOneIter()來建樹


2.learner

上一節,在learner中初始化目標函數和模型,主要賦值給了這兩個變量。

/include/xgboost/learner.h

  /*! \brief objective function */
  std::unique_ptr<ObjFunction> obj_;
  /*! \brief The gradient booster used by the model*/
  std::unique_ptr<GradientBooster> gbm_;

ObjFunction是基類,定義了很多虛函數。類RegLossObj等都繼承於此類,主要實現了根據不同的模型和loss,將一階二階導數計算出來。


以線性迴歸模型爲例

src\objective\regression_obj.cc

定義平方損失函數:

// linear regression
struct LinearSquareLoss {
  static bst_float PredTransform(bst_float x) { return x; }
  static bool CheckLabel(bst_float x) { return true; }
  static bst_float FirstOrderGradient(bst_float predt, bst_float label) { return predt - label; }
  static bst_float SecondOrderGradient(bst_float predt, bst_float label) { return 1.0f; }
  static bst_float ProbToMargin(bst_float base_score) { return base_score; }
  static const char* LabelErrorMsg() { return ""; }
  static const char* DefaultEvalMetric() { return "rmse"; }
}


class RegLossObj : public ObjFunction{
...
void GetGradient(const std::vector<bst_float> &preds,
                   const MetaInfo &info,
                   int iter,
                   std::vector<bst_gpair> *out_gpair) override {
...
    out_gpair->resize(preds.size());
    // check if label in range
    bool label_correct = true;
    // start calculating gradient
    const omp_ulong ndata = static_cast<omp_ulong>(preds.size());
    #pragma omp parallel for schedule(static)
    for (omp_ulong i = 0; i < ndata; ++i) {
      bst_float p = Loss::PredTransform(preds[i]);
      bst_float w = info.GetWeight(i);
      if (info.labels[i] == 1.0f) w *= param_.scale_pos_weight;
      if (!Loss::CheckLabel(info.labels[i])) label_correct = false;
      out_gpair->at(i) = bst_gpair(Loss::FirstOrderGradient(p, info.labels[i]) * w,
                                   Loss::SecondOrderGradient(p, info.labels[i]) * w);
    
    }
...
}
可以發現在計算一階導和二階導的時候,採用了並行處理。以一條數據作爲一個粒度。


GradientBooster是基類,有兩個模型繼承於此類。XGBoost中除了有Tree模型(GBTree),同時也實現了線性模型(GBLinear)。與梯度下降和牛頓法不同的是,在每次迭代的過程中,每個屬性單獨計算,採用類似於一維的牛頓法來更新一個屬性。這裏不是重點,有興趣的同學可以看\src\gbm\gblinear.cc。


3.UpdateOneIter


這裏可以看出每次迭代的操作,主要有:

  void UpdateOneIter(int iter, DMatrix* train) override {
  	
    this->LazyInitDMatrix(train);
    this->PredictRaw(train, &preds_);                                   //獲取上一輪預測值
    obj_->GetGradient(preds_, train->info(), iter, &gpair_);            //計算一階導和二階導
    gbm_->DoBoost(train, &gpair_, obj_.get());                          //建樹
  }

在DoBoost 中確定了輸出維度,調用BoostNewTrees構造森林。最後終於調用了update在輸出維度不爲1時,調用BoostNewTrees是並行的。這裏構造不止一棵樹是爲了一下一步隨機森林做準備。因爲每一次迭代,是分給了多個樹完成,而不是一顆。DoBoost這部分代碼可以在gbtree.cc中找到,這裏不再贅述。


4.ColMaker

4.1 用於統計的結構體

ColMaker繼承於TreeUpdater,是單機多線程建樹的實現過程,也是本文重點記錄的地方。

ColMaker提供兩個函數:Init、Update,分別用於初始化和更新建樹由上層調用,同時爲了實現多線程,定義了三個結構體:ThreadEntry、NodeEntry、Builder。

class ColMaker: public TreeUpdater {
 public:
       void Init(const std::vector<std::pair<std::string, std::string> >& args) override
       void Update(const std::vector<bst_gpair> &gpair,DMatrix* dmat,const std::vector<RegTree*> &trees) override
	// training parameter
       TrainParam param;
        // data structure
       /*! \brief per thread x per node entry to store tmp data */
       struct ThreadEntr
       struct NodeEntry
        // actual builder that runs the algorithm
       struct Builder  
}


首先看一下結構體ThreadEntry

struct ThreadEntry {
    /*! \brief statistics of data */
    TStats stats;
    /*! \brief extra statistics of data */
    TStats stats_extra;
    /*! \brief last feature value scanned */
    bst_float last_fvalue;
    /*! \brief first feature value scanned */
    bst_float first_fvalue;
    /*! \brief current best solution */
    SplitEntry best;
    // constructor
    explicit ThreadEntry(const TrainParam &param)
        : stats(param), stats_extra(param) {
    }

ThreadEntry結構體用於多線程中單個線程內的樣本統計,一個線程擁有一個ThreadEntry 變量。


struct NodeEntry {
    /*! \brief statics for node entry */
    TStats stats;
    /*! \brief loss of this node, without split */
    bst_float root_gain;
    /*! \brief weight calculated related to current data */
    bst_float weight;
    /*! \brief current best solution */
    SplitEntry best;
    // constructor
    explicit NodeEntry(const TrainParam& param)
        : stats(param), root_gain(0.0f), weight(0.0f){
    }

NodeEntry 用於一個樹節點,所擁有的樣本的統計信息。一個節點在分列前,需要將自己的樣本分給不同的線程進行處理,最後得到多個ThreadEntry 變量,彙總到每個節點自己的NodeEntry 變量中。詳細過程如下。


4.2 建樹總體流程

Update -> Builder.Update


struct Builder {
...
   public:
    // constructor
    explicit Builder(const TrainParam& param) : param(param), nthread(omp_get_max_threads()) {}
    // update one tree, growing
    virtual void Update(const std::vector<bst_gpair>& gpair,
                        DMatrix* p_fmat,
                        RegTree* p_tree) {
      this->InitData(gpair, *p_fmat, *p_tree);
      this->InitNewNode(qexpand_, gpair, *p_fmat, *p_tree);
      for (int depth = 0; depth < param.max_depth; ++depth) {
        this->FindSplit(depth, qexpand_, gpair, p_fmat, p_tree);
        this->ResetPosition(qexpand_, p_fmat, *p_tree);
        this->UpdateQueueExpand(*p_tree, &qexpand_);
        this->InitNewNode(qexpand_, gpair, *p_fmat, *p_tree);
        // if nothing left to be expand, break
        if (qexpand_.size() == 0) break;
      }
      // set all the rest expanding nodes to leaf
      for (size_t i = 0; i < qexpand_.size(); ++i) {
        const int nid = qexpand_[i];
        (*p_tree)[nid].set_leaf(snode[nid].weight * param.learning_rate);
      }
      // remember auxiliary statistics in the tree node
      for (int nid = 0; nid < p_tree->param.num_nodes; ++nid) {
        p_tree->stat(nid).loss_chg = snode[nid].best.loss_chg;
        p_tree->stat(nid).base_weight = snode[nid].weight;
        p_tree->stat(nid).sum_hess = static_cast<float>(snode[nid].stats.sum_hess);
        snode[nid].stats.SetLeafVec(param, p_tree->leafvec(nid));
      }
    }
...
}



上述代碼的流程可以描述爲:

初始化變量
初始化節點
For (depth 1:max_depth){
       找分割點,更新變量p_tree,維護樹的增長
       重置樣本所屬節點,更新變量position,主要將樣本從父節點分配到剛分裂出來的葉子節點中
       將葉子預分裂的節點挑出來,更新變量qexpand_,這個變量是爲了存儲待處理的新節點
       初始化新節點
}

qexpand_中的節點設置爲葉子節點
把隊列裏的節點信息加到樹模型中


4.3 抽樣

在XGBoost中,爲了防止過擬合,有很多地方用到了抽樣。其中一種是通過參數對樣本進行抽樣:

Builder -> Update - >InitData

 // mark subsample
        if (param.subsample < 1.0f) {
          std::bernoulli_distribution coin_flip(param.subsample);
          auto& rnd = common::GlobalRandom();
          for (size_t i = 0; i < rowset.size(); ++i) {
            const bst_uint ridx = rowset[i];
            if (gpair[ridx].hess < 0.0f) continue;
            if (!coin_flip(rnd)) position[ridx] = ~position[ridx];
          }

在XGBoost的設計中,position[ridx]如果小於0,則認爲該樣本被刪除。由於loss是凸函數,二階導必定大於0,所以如果小於0,可能在計算過程中出現溢出等異常出現,則不考慮該樣本。所以如果不將樣本加入到計算中,只需取反。在上述代碼中,是通過伯努利分佈生成隨機數來進行抽樣。


XGBoost還對模型的屬性進行抽樣:

Builder -> Update - > FindSplit

std::vector<bst_uint> feat_set = feat_index;
if (param.colsample_bylevel != 1.0f) {
      std::shuffle(feat_set.begin(), feat_set.end(), common::GlobalRandom());
      unsigned n = static_cast<unsigned>(param.colsample_bylevel * feat_index.size());
      feat_set.resize(n);
}

利用std::shuffle隨機排序取前n來實現列抽樣

4.4 尋找分割點

Builder -> Update - > FindSplit 

    // find splits at current level, do split per level
    inline void FindSplit(int depth,
                          const std::vector<int> &qexpand,
                          const std::vector<bst_gpair> &gpair,
                          DMatrix *p_fmat,
                          RegTree *p_tree) {
      ...
      dmlc::DataIter<ColBatch>* iter = p_fmat->ColIterator(feat_set);
      while (iter->Next()) {
        this->UpdateSolution(iter->Value(), gpair, *p_fmat);
      }
      // after this each thread's stemp will get the best candidates, aggregate results
      this->SyncBestSolution(qexpand);
      // get the best result, we can synchronize the solution
      for (size_t i = 0; i < qexpand.size(); ++i) {
        const int nid = qexpand[i];
        NodeEntry &e = snode[nid];
        // now we know the solution in snode[nid], set split
        if (e.best.loss_chg > rt_eps) {
          p_tree->AddChilds(nid);
          (*p_tree)[nid].set_split(e.best.split_index(), e.best.split_value, e.best.default_left());
          // mark right child as 0, to indicate fresh leaf
          (*p_tree)[(*p_tree)[nid].cleft()].set_leaf(0.0f, 0);
          (*p_tree)[(*p_tree)[nid].cright()].set_leaf(0.0f, 0);
        } else {
          (*p_tree)[nid].set_leaf(e.weight * param.learning_rate);
        }
      }
    }

XGBoost把樣本轉換成了以列爲基礎的迭代器。這一部分設計等我整理一下再傳上來。代碼將數據以列爲基礎通過調用UpdateSolution計算分割點。爲什麼要這麼做呢,這裏已經能想到了,每一列至少交給了一個線程去處理。而SyncBestSolution則是將線程處理後的數據進行彙總的過程。

而進入UpdateSolution後,如果數據量不大則按屬性進行多線程處理了,每個線程通過調用EnumerateSplit方法按實現。但是當數據量比較大,或者通過參數設置parallel_option,可以在此基礎上進行更細的多線程拆分。

通過函數ParallelFindSplit實現:

Builder -> Update - > FindSplit -> ParallelFindSplit

    // parallel find the best split of current fid
    // this function does not support nested functions
    inline void ParallelFindSplit(const ColBatch::Inst &col,
                                  bst_uint fid,
                                  const DMatrix &fmat,
                                  const std::vector<bst_gpair> &gpair) {
      // TODO(tqchen): double check stats order.
      const MetaInfo& info = fmat.info();
      const bool ind = col.length != 0 && col.data[0].fvalue == col.data[col.length - 1].fvalue;
      bool need_forward = param.need_forward_search(fmat.GetColDensity(fid), ind);
      bool need_backward = param.need_backward_search(fmat.GetColDensity(fid), ind);
      const std::vector<int> &qexpand = qexpand_;
      #pragma omp parallel
      {
        const int tid = omp_get_thread_num();
        std::vector<ThreadEntry> &temp = stemp[tid];
        // cleanup temp statistics
        for (size_t j = 0; j < qexpand.size(); ++j) {
          temp[qexpand[j]].stats.Clear();
        }
        bst_uint step = (col.length + this->nthread - 1) / this->nthread;
        bst_uint end = std::min(col.length, step * (tid + 1));
        for (bst_uint i = tid * step; i < end; ++i) {
          const bst_uint ridx = col[i].index;
          const int nid = position[ridx];
          if (nid < 0) continue;
          const bst_float fvalue = col[i].fvalue;
          if (temp[nid].stats.Empty()) {
            temp[nid].first_fvalue = fvalue;
          }
          temp[nid].stats.Add(gpair, info, ridx);
          temp[nid].last_fvalue = fvalue;
        }
      }
      // start collecting the partial sum statistics
      bst_omp_uint nnode = static_cast<bst_omp_uint>(qexpand.size());
      #pragma omp parallel for schedule(static)
      for (bst_omp_uint j = 0; j < nnode; ++j) {
        const int nid = qexpand[j];
        TStats sum(param), tmp(param), c(param);
        for (int tid = 0; tid < this->nthread; ++tid) {
          tmp = stemp[tid][nid].stats;
          stemp[tid][nid].stats = sum;
          sum.Add(tmp);
          if (tid != 0) {
            std::swap(stemp[tid - 1][nid].last_fvalue, stemp[tid][nid].first_fvalue);
          }
        }
        for (int tid = 0; tid < this->nthread; ++tid) {
          stemp[tid][nid].stats_extra = sum;
          ThreadEntry &e = stemp[tid][nid];
          bst_float fsplit;
          if (tid != 0) {
            if (stemp[tid - 1][nid].last_fvalue != e.first_fvalue) {
              fsplit = (stemp[tid - 1][nid].last_fvalue + e.first_fvalue) * 0.5f;
            } else {
              continue;
            }
          } else {
            fsplit = e.first_fvalue - rt_eps;
          }
          if (need_forward && tid != 0) {
            c.SetSubstract(snode[nid].stats, e.stats);
            if (c.sum_hess >= param.min_child_weight &&
                e.stats.sum_hess >= param.min_child_weight) {
              bst_float loss_chg = static_cast<bst_float>(
                  constraints_[nid].CalcSplitGain(param, fid, e.stats, c) - snode[nid].root_gain);
              e.best.Update(loss_chg, fid, fsplit, false);
            }
          }
          if (need_backward) {
            tmp.SetSubstract(sum, e.stats);
            c.SetSubstract(snode[nid].stats, tmp);
            if (c.sum_hess >= param.min_child_weight &&
                tmp.sum_hess >= param.min_child_weight) {
              bst_float loss_chg = static_cast<bst_float>(
                  constraints_[nid].CalcSplitGain(param, fid, tmp, c) - snode[nid].root_gain);
              e.best.Update(loss_chg, fid, fsplit, true);
            }
          }
        }
        if (need_backward) {
          tmp = sum;
          ThreadEntry &e = stemp[this->nthread-1][nid];
          c.SetSubstract(snode[nid].stats, tmp);
          if (c.sum_hess >= param.min_child_weight &&
              tmp.sum_hess >= param.min_child_weight) {
            bst_float loss_chg = static_cast<bst_float>(
                constraints_[nid].CalcSplitGain(param, fid, tmp, c) - snode[nid].root_gain);
            e.best.Update(loss_chg, fid, e.last_fvalue + rt_eps, true);
          }
        }
      }
      // rescan, generate candidate split
      #pragma omp parallel
      {
        TStats c(param), cright(param);
        const int tid = omp_get_thread_num();
        std::vector<ThreadEntry> &temp = stemp[tid];
        bst_uint step = (col.length + this->nthread - 1) / this->nthread;
        bst_uint end = std::min(col.length, step * (tid + 1));
        for (bst_uint i = tid * step; i < end; ++i) {
          const bst_uint ridx = col[i].index;
          const int nid = position[ridx];
          if (nid < 0) continue;
          const bst_float fvalue = col[i].fvalue;
          // get the statistics of nid
          ThreadEntry &e = temp[nid];
          if (e.stats.Empty()) {
            e.stats.Add(gpair, info, ridx);
            e.first_fvalue = fvalue;
          } else {
            // forward default right
            if (fvalue != e.first_fvalue) {
              if (need_forward) {
                c.SetSubstract(snode[nid].stats, e.stats);
                if (c.sum_hess >= param.min_child_weight &&
                    e.stats.sum_hess >= param.min_child_weight) {
                  bst_float loss_chg = static_cast<bst_float>(
                      constraints_[nid].CalcSplitGain(param, fid, e.stats, c) -
                      snode[nid].root_gain);
                  e.best.Update(loss_chg, fid, (fvalue + e.first_fvalue) * 0.5f, false);
                }
              }
              if (need_backward) {
                cright.SetSubstract(e.stats_extra, e.stats);
                c.SetSubstract(snode[nid].stats, cright);
                if (c.sum_hess >= param.min_child_weight &&
                    cright.sum_hess >= param.min_child_weight) {
                  bst_float loss_chg = static_cast<bst_float>(
                      constraints_[nid].CalcSplitGain(param, fid, c, cright) -
                      snode[nid].root_gain);
                  e.best.Update(loss_chg, fid, (fvalue + e.first_fvalue) * 0.5f, true);
                }
              }
            }
            e.stats.Add(gpair, info, ridx);
            e.first_fvalue = fvalue;
          }
        }
      }
    }


1. 上述代碼將某一屬性的樣本空間的數據進行分塊,塊大小爲(col.length + this->nthread - 1) / this->nthread),這一步是並行的,每一塊是一個線程,每個線程統計屬於自己的樣本的每個節點的一階導和二階導數的和,還需要統計邊緣值(first_fvalue、last_fvalue)

2. 按照塊邊緣,初始化線程變量stemp[tid][nid].best,這一步是節點並行,因爲需要用到兩個相鄰線程所用有的數據塊的邊緣值:stemp[tid - 1][nid].last_fvalue + e.first_fvalue) * 0.5f,所以無法按照塊並行

3. 第三步則是按照數據塊並行,每個線程處理自己的數據塊來更新stemp[tid][nid].best,這樣每個線程都有一個根據自己樣本計算出來的屬於自己的最優分割點。

4. 最後再根據上文中提到的SyncBestSolution來彙總所有線程,最後進行樹的增長、qexpand更新、position更新等操作。

5. 由於迭代器是按照屬性排序的,e.stats.Add負責按順序累加樣本,再用總的減去就是剩下的:c.SetSubstract(snode[nid].stats, tmp),所以計算分割點只需要遍歷一遍樣本即可。

6. constraints_的類型是模板類傳遞進來的,具體分爲ValueConstraint和NoConstraint。ValueConstraint具有上下界約束,具體定義參看src\tree\param.h



總結:

本文梳理了XGBoost 在訓練過程的源碼的流程,通過閱讀,瞭解了XGBoost的運行機制、設計模式爲以後自己編寫代碼提供了寶貴經驗。




發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章