Implement vector leaf for random forest

cpp/bench/sg/fil.cu

@@ -81,7 +81,7 @@ class FIL : public RegressionFixture<float> {
     fit(*handle, mPtr, data.X.data(), train_nrows, params.ncols, data.y.data(), p_rest.rf);
     CUDA_CHECK(cudaStreamSynchronize(stream));
 
-    ML::build_treelite_forest(&model, &rf_model, params.ncols, params.nclasses > 1 ? 2 : 1);
+    ML::build_treelite_forest(&model, &rf_model, params.ncols);
     ML::fil::treelite_params_t tl_params = {
       .algo              = p_rest.algo,
       .output_class      = params.nclasses > 1,    // cuML RF forest

cpp/include/cuml/ensemble/randomforest.hpp

@@ -130,8 +130,7 @@ std::string get_rf_json(const RandomForestMetaData<T, L>* forest);
 template <class T, class L>
 void build_treelite_forest(ModelHandle* model,
                            const RandomForestMetaData<T, L>* forest,
-                           int num_features,
-                           int task_category);
+                           int num_features);
 
 ModelHandle concatenate_trees(std::vector<ModelHandle> treelite_handles);
 

cpp/include/cuml/tree/decisiontree.hpp

@@ -103,7 +103,9 @@ struct TreeMetaDataNode {
   int depth_counter;
   int leaf_counter;
   double train_time;
+  std::vector<T> vector_leaf;
   std::vector<SparseTreeNode<T, L>> sparsetree;
+  int num_outputs;
 };
 
 /**

cpp/include/cuml/tree/flatnode.h

@@ -33,20 +33,14 @@
 template <typename DataT, typename LabelT, typename IdxT = int>
 struct SparseTreeNode {
  private:
-  LabelT prediction     = LabelT(0);
   IdxT colid            = 0;
   DataT quesval         = DataT(0);
   DataT best_metric_val = DataT(0);
   IdxT left_child_id    = -1;
   IdxT instance_count   = 0;
-  FLATNODE_HD SparseTreeNode(LabelT prediction,
-                             IdxT colid,
-                             DataT quesval,
-                             DataT best_metric_val,
-                             int64_t left_child_id,
-                             IdxT instance_count)
-    : prediction(prediction),
-      colid(colid),
+  FLATNODE_HD SparseTreeNode(
+    IdxT colid, DataT quesval, DataT best_metric_val, int64_t left_child_id, IdxT instance_count)
+    : colid(colid),
       quesval(quesval),
       best_metric_val(best_metric_val),
       left_child_id(left_child_id),
@@ -55,7 +49,6 @@ struct SparseTreeNode {
   }
 
  public:
-  FLATNODE_HD LabelT Prediction() const { return prediction; }
   FLATNODE_HD IdxT ColumnId() const { return colid; }
   FLATNODE_HD DataT QueryValue() const { return quesval; }
   FLATNODE_HD DataT BestMetric() const { return best_metric_val; }
@@ -67,17 +60,17 @@ struct SparseTreeNode {
     IdxT colid, DataT quesval, DataT best_metric_val, int64_t left_child_id, IdxT instance_count)
   {
     return SparseTreeNode<DataT, LabelT>{
-      LabelT(0), colid, quesval, best_metric_val, left_child_id, instance_count};
+      colid, quesval, best_metric_val, left_child_id, instance_count};
   }
-  FLATNODE_HD static SparseTreeNode CreateLeafNode(LabelT prediction, IdxT instance_count)
+  FLATNODE_HD static SparseTreeNode CreateLeafNode(IdxT instance_count)
   {
-    return SparseTreeNode<DataT, LabelT>{prediction, 0, 0, 0, -1, instance_count};
+    return SparseTreeNode<DataT, LabelT>{0, 0, 0, -1, instance_count};
   }
   FLATNODE_HD bool IsLeaf() const { return left_child_id == -1; }
   bool operator==(const SparseTreeNode& other) const
   {
-    return (this->prediction == other.prediction) && (this->colid == other.colid) &&
-           (this->quesval == other.quesval) && (this->best_metric_val == other.best_metric_val) &&
+    return (this->colid == other.colid) && (this->quesval == other.quesval) &&
+           (this->best_metric_val == other.best_metric_val) &&
            (this->left_child_id == other.left_child_id) &&
            (this->instance_count == other.instance_count);
   }

cpp/src/decisiontree/batched-levelalgo/builder.cuh

@@ -50,11 +50,12 @@ class NodeQueue {
   std::deque<NodeWorkItem> work_items_;
 
  public:
-  NodeQueue(DecisionTreeParams params, size_t max_nodes, size_t sampled_rows)
+  NodeQueue(DecisionTreeParams params, size_t max_nodes, size_t sampled_rows, int num_outputs)
     : params(params), tree(std::make_shared<DT::TreeMetaDataNode<DataT, LabelT>>())
   {
+    tree->num_outputs = num_outputs;
     tree->sparsetree.reserve(max_nodes);
-    tree->sparsetree.emplace_back(NodeT::CreateLeafNode(0, sampled_rows));
+    tree->sparsetree.emplace_back(NodeT::CreateLeafNode(sampled_rows));
     tree->leaf_counter  = 1;
     tree->depth_counter = 0;
     node_instances_.reserve(max_nodes);
@@ -112,7 +113,7 @@ class NodeQueue {
                                                              parent_range.count);
       tree->leaf_counter++;
       // left
-      tree->sparsetree.emplace_back(NodeT::CreateLeafNode(0, split.nLeft));
+      tree->sparsetree.emplace_back(NodeT::CreateLeafNode(split.nLeft));
       node_instances_.emplace_back(InstanceRange{parent_range.begin, std::size_t(split.nLeft)});
 
       // Do not add a work item if this child is definitely a leaf
@@ -122,7 +123,7 @@ class NodeQueue {
       }
 
       // right
-      tree->sparsetree.emplace_back(NodeT::CreateLeafNode(0, parent_range.count - split.nLeft));
+      tree->sparsetree.emplace_back(NodeT::CreateLeafNode(parent_range.count - split.nLeft));
       node_instances_.emplace_back(
         InstanceRange{parent_range.begin + split.nLeft, parent_range.count - split.nLeft});
 
@@ -314,14 +315,14 @@ struct Builder {
   {
     ML::PUSH_RANGE("Builder::train @builder.cuh [batched-levelalgo]");
     MLCommon::TimerCPU timer;
-    NodeQueue<DataT, LabelT> queue(params, this->maxNodes(), input.nSampledRows);
+    NodeQueue<DataT, LabelT> queue(params, this->maxNodes(), input.nSampledRows, input.numOutputs);
     while (queue.HasWork()) {
       auto work_items                      = queue.Pop();
       auto [splits_host_ptr, splits_count] = doSplit(work_items);
       queue.Push(work_items, splits_host_ptr);
     }
     auto tree = queue.GetTree();
-    this->SetLeafPredictions(&tree->sparsetree, queue.GetInstanceRanges());
+    this->SetLeafPredictions(tree, queue.GetInstanceRanges());
     tree->train_time = timer.getElapsedMilliseconds();
     ML::POP_RANGE();
     return tree;
@@ -443,30 +444,41 @@ struct Builder {
   }
 
   // Set the leaf value predictions in batch
-  void SetLeafPredictions(std::vector<NodeT>* tree,
+  void SetLeafPredictions(std::shared_ptr<DT::TreeMetaDataNode<DataT, LabelT>> tree,
                           const std::vector<InstanceRange>& instance_ranges)
   {
+    tree->vector_leaf.resize(tree->sparsetree.size() * input.numOutputs);
+    ASSERT(tree->sparsetree.size() == instance_ranges.size(),
+           "Expected instance range for each node");
     // do this in batch to reduce peak memory usage in extreme cases
-    std::size_t max_batch_size = min(std::size_t(100000), tree->size());
+    std::size_t max_batch_size = min(std::size_t(100000), tree->sparsetree.size());
     rmm::device_uvector<NodeT> d_tree(max_batch_size, handle.get_stream());
     rmm::device_uvector<InstanceRange> d_instance_ranges(max_batch_size, handle.get_stream());
+    rmm::device_uvector<DataT> d_leaves(max_batch_size * input.numOutputs, handle.get_stream());
+
     ObjectiveT objective(input.numOutputs, params.min_impurity_decrease, params.min_samples_leaf);
-    for (std::size_t batch_begin = 0; batch_begin < tree->size(); batch_begin += max_batch_size) {
-      std::size_t batch_end  = min(batch_begin + max_batch_size, tree->size());
+    for (std::size_t batch_begin = 0; batch_begin < tree->sparsetree.size();
+         batch_begin += max_batch_size) {
+      std::size_t batch_end  = min(batch_begin + max_batch_size, tree->sparsetree.size());
       std::size_t batch_size = batch_end - batch_begin;
       raft::update_device(
-        d_tree.data(), tree->data() + batch_begin, batch_size, handle.get_stream());
+        d_tree.data(), tree->sparsetree.data() + batch_begin, batch_size, handle.get_stream());
       raft::update_device(d_instance_ranges.data(),
                           instance_ranges.data() + batch_begin,
                           batch_size,
                           handle.get_stream());
+
+      CUDA_CHECK(
+        cudaMemsetAsync(d_leaves.data(), 0, sizeof(DataT) * d_leaves.size(), handle.get_stream()));
       size_t smemSize = sizeof(BinT) * input.numOutputs;
       int num_blocks  = batch_size;
       leafKernel<<<num_blocks, TPB_DEFAULT, smemSize, handle.get_stream()>>>(
-        objective, input, d_tree.data(), d_instance_ranges.data());
-      raft::update_host(tree->data() + batch_begin, d_tree.data(), batch_size, handle.get_stream());
+        objective, input, d_tree.data(), d_instance_ranges.data(), d_leaves.data());
+      raft::update_host(tree->vector_leaf.data() + batch_begin * input.numOutputs,
+                        d_leaves.data(),
+                        batch_size * input.numOutputs,
+                        handle.get_stream());
     }
-    CUDA_CHECK(cudaStreamSynchronize(handle.get_stream()));
   }
 };  // end Builder
 

cpp/src/decisiontree/batched-levelalgo/kernels.cuh

@@ -123,7 +123,6 @@ DI void partitionSamples(const Input<DataT, LabelT, IdxT>& input,
     }
   }
 }
-
 template <typename DataT, typename LabelT, typename IdxT, typename ObjectiveT, int TPB>
 __global__ void nodeSplitKernel(IdxT max_depth,
                                 IdxT min_samples_leaf,
@@ -144,17 +143,19 @@ __global__ void nodeSplitKernel(IdxT max_depth,
   partitionSamples<DataT, LabelT, IdxT, TPB>(input, split, work_item, (char*)smem);
 }
 
-template <typename InputT, typename NodeT, typename ObjectiveT>
+template <typename InputT, typename NodeT, typename ObjectiveT, typename DataT>
 __global__ void leafKernel(ObjectiveT objective,
                            InputT input,
-                           NodeT* tree,
-                           const InstanceRange* instance_ranges)
+                           const NodeT* tree,
+                           const InstanceRange* instance_ranges,
+                           DataT* leaves)
 {
   using BinT = typename ObjectiveT::BinT;
   extern __shared__ char shared_memory[];
   auto histogram = reinterpret_cast<BinT*>(shared_memory);
-  auto& node     = tree[blockIdx.x];
-  auto range     = instance_ranges[blockIdx.x];
+  auto node_id   = blockIdx.x;
+  auto& node     = tree[node_id];
+  auto range     = instance_ranges[node_id];
   if (!node.IsLeaf()) return;
   auto tid = threadIdx.x;
   for (int i = tid; i < input.numOutputs; i += blockDim.x) {
@@ -167,8 +168,7 @@ __global__ void leafKernel(ObjectiveT objective,
   }
   __syncthreads();
   if (tid == 0) {
-    node =
-      NodeT::CreateLeafNode(ObjectiveT::LeafPrediction(histogram, input.numOutputs), range.count);
+    ObjectiveT::SetLeafVector(histogram, input.numOutputs, leaves + input.numOutputs * node_id);
   }
 }
 
@@ -338,8 +338,8 @@ __global__ void computeSplitKernel(BinT* hist,
   if (input.nSampledCols == input.N) {
     col = colStart + blockIdx.y;
   } else {
-    int colIndex = colStart + blockIdx.y;
-    col          = select(colIndex, treeid, work_item.idx, seed, input.N);
+    IdxT colIndex = colStart + blockIdx.y;
+    col           = select(colIndex, treeid, work_item.idx, seed, input.N);
   }
 
   // populating shared memory with initial values

cpp/src/decisiontree/batched-levelalgo/metrics.cuh

@@ -73,7 +73,7 @@ class GiniObjectiveFunction {
     constexpr DataT One = DataT(1.0);
     DataT invlen        = One / len;
     for (IdxT i = threadIdx.x; i < nbins; i += blockDim.x) {
-      int nLeft = 0;
+      IdxT nLeft = 0;
       for (IdxT j = 0; j < nclasses; ++j) {
         nLeft += scdf_labels[nbins * j + i].x;
       }
@@ -106,18 +106,16 @@ class GiniObjectiveFunction {
     }
     return sp;
   }
-  static DI LabelT LeafPrediction(BinT* shist, int nclasses)
+  static DI void SetLeafVector(BinT* shist, int nclasses, DataT* out)
   {
-    int class_idx = 0;
-    int count     = 0;
+    // Output probability
+    int total = 0;
     for (int i = 0; i < nclasses; i++) {
-      auto current_count = shist[i].x;
-      if (current_count > count) {
-        class_idx = i;
-        count     = current_count;
-      }
+      total += shist[i].x;
+    }
+    for (int i = 0; i < nclasses; i++) {
+      out[i] = DataT(shist[i].x) / total;
     }
-    return class_idx;
   }
 };
 
@@ -146,7 +144,7 @@ class EntropyObjectiveFunction {
     constexpr DataT One = DataT(1.0);
     DataT invlen        = One / len;
     for (IdxT i = threadIdx.x; i < nbins; i += blockDim.x) {
-      int nLeft = 0;
+      IdxT nLeft = 0;
       for (IdxT j = 0; j < nclasses; ++j) {
         nLeft += scdf_labels[nbins * j + i].x;
       }
@@ -185,10 +183,16 @@ class EntropyObjectiveFunction {
     }
     return sp;
   }
-  static DI LabelT LeafPrediction(BinT* shist, int nclasses)
+  static DI void SetLeafVector(BinT* shist, int nclasses, DataT* out)
   {
-    // Same as Gini
-    return GiniObjectiveFunction<DataT, LabelT, IdxT>::LeafPrediction(shist, nclasses);
+    // Output probability
+    int total = 0;
+    for (int i = 0; i < nclasses; i++) {
+      total += shist[i].x;
+    }
+    for (int i = 0; i < nclasses; i++) {
+      out[i] = DataT(shist[i].x) / total;
+    }
   }
 };
 
@@ -260,9 +264,11 @@ class MSEObjectiveFunction {
     return sp;
   }
 
-  static DI LabelT LeafPrediction(BinT* shist, int nclasses)
+  static DI void SetLeafVector(BinT* shist, int nclasses, DataT* out)
   {
-    return shist[0].label_sum / shist[0].count;
+    for (int i = 0; i < nclasses; i++) {
+      out[i] = shist[i].label_sum / shist[i].count;
+    }
   }
 };
 

cpp/src/decisiontree/decisiontree.cu

@@ -95,14 +95,14 @@ template <class T, class L>
 std::string get_tree_text(const TreeMetaDataNode<T, L>* tree)
 {
   std::string summary = get_tree_summary_text<T, L>(tree);
-  return summary + "\n" + get_node_text<T, L>("", tree->sparsetree, 0, false);
+  return summary + "\n" + get_node_text<T, L>("", tree, 0, false);
 }
 
 template <class T, class L>
 std::string get_tree_json(const TreeMetaDataNode<T, L>* tree)
 {
   std::ostringstream oss;
-  return get_node_json("", tree->sparsetree, 0);
+  return get_node_json("", tree, 0);
 }
 
 // Functions' specializations