Support XGBoost classifiers with num_class > 2, num_parallel_tree > 1

src/frontend/xgboost.cc

@@ -380,8 +380,10 @@ inline std::unique_ptr<treelite::Model> ParseStream(std::istream& fi) {
     xgb_trees_.emplace_back();
     xgb_trees_.back().Load(fp.get());
   }
-  TREELITE_CHECK_EQ(gbm_param_.num_roots, 1) << "multi-root trees not supported";
-  // tree_info is currently unused.
+  if (mparam_.major_version < 1 || (mparam_.major_version == 1 && mparam_.minor_version < 6)) {
+    // In XGBoost 1.6, num_roots is used as num_parallel_tree, so don't check
+    TREELITE_CHECK_EQ(gbm_param_.num_roots, 1) << "multi-root trees not supported";
+  }
   std::vector<int> tree_info;
   tree_info.resize(gbm_param_.num_trees);
   if (gbm_param_.num_trees > 0) {
@@ -468,6 +470,36 @@ inline std::unique_ptr<treelite::Model> ParseStream(std::istream& fi) {
       tree.SetSumHess(new_id, stat.sum_hess);
     }
   }
+
+  // Special handling for multi-class classifier with num_parallel_tree > 1
+  if (num_class > 1) {
+    // Infer num_parallel_tree
+    unsigned num_parallel_tree = 0;
+    for (int e : tree_info) {
+      if (e != 0) {
+        break;
+      }
+      ++num_parallel_tree;
+    }
+    if (num_parallel_tree > 1) {
+      // Re-order trees to recover the grove-per-class layout.
+      // The prediction for the i-th class is determined by the trees whose index is congruent
+      // to [i] modulo [num_class].
+      // Currently, the trees' association with classes is as follows:
+      // 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2
+      // We need to re-order them as follows:
+      // 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2
+      std::vector<treelite::Tree<float, float>> new_trees;
+      std::size_t num_tree = model->trees.size();
+      for (std::size_t c = 0; c < num_parallel_tree; ++c) {
+        for (std::size_t tree_id = c; tree_id < num_tree; tree_id += num_parallel_tree) {
+          new_trees.push_back(std::move(model->trees[tree_id]));
+        }
+      }
+      TREELITE_CHECK_EQ(new_trees.size(), num_tree);
+      model->trees = std::move(new_trees);
+    }
+  }
   return model_ptr;
 }
 

src/frontend/xgboost/xgboost_json.h

@@ -227,6 +227,14 @@ class ArrayHandler : public OutputHandler<std::vector<ElemType>> {
   }
 };
 
+struct ParsedXGBoostModel {
+  std::unique_ptr<treelite::Model> model_ptr;
+  treelite::ModelImpl<float, float>* model;
+  std::vector<unsigned> version;
+  std::vector<int> tree_info;
+  std::string objective_name;
+};
+
 /*! \brief handler for TreeParam objects from XGBoost schema*/
 class TreeParamHandler : public OutputHandler<int> {
  public:
@@ -266,16 +274,16 @@ class RegTreeHandler : public OutputHandler<treelite::Tree<float, float>> {
 };
 
 /*! \brief handler for GBTreeModel objects from XGBoost schema*/
-class GBTreeModelHandler : public OutputHandler<treelite::ModelImpl<float, float>> {
-  using OutputHandler<treelite::ModelImpl<float, float>>::OutputHandler;
+class GBTreeModelHandler : public OutputHandler<ParsedXGBoostModel> {
+  using OutputHandler<ParsedXGBoostModel>::OutputHandler;
   bool StartArray() override;
   bool StartObject() override;
 };
 
 /*! \brief handler for GradientBoosterHandler objects from XGBoost schema*/
-class GradientBoosterHandler : public OutputHandler<treelite::ModelImpl<float, float>> {
+class GradientBoosterHandler : public OutputHandler<ParsedXGBoostModel> {
  public:
-  using OutputHandler<treelite::ModelImpl<float, float>>::OutputHandler;
+  using OutputHandler<ParsedXGBoostModel>::OutputHandler;
   bool String(const char *str, std::size_t length, bool copy) override;
   bool StartArray() override;
   bool StartObject() override;
@@ -301,16 +309,10 @@ class LearnerParamHandler : public OutputHandler<treelite::ModelImpl<float, floa
   bool String(const char *str, std::size_t length, bool copy) override;
 };
 
-struct XGBoostModelHandle {
-  treelite::ModelImpl<float, float>* model;
-  std::vector<unsigned> version;
-  std::string objective_name;
-};
-
 /*! \brief handler for Learner objects from XGBoost schema*/
-class LearnerHandler : public OutputHandler<XGBoostModelHandle> {
+class LearnerHandler : public OutputHandler<ParsedXGBoostModel> {
  public:
-  using OutputHandler<XGBoostModelHandle>::OutputHandler;
+  using OutputHandler<ParsedXGBoostModel>::OutputHandler;
   bool StartObject() override;
   bool EndObject(std::size_t memberCount) override;
   bool StartArray() override;
@@ -320,29 +322,27 @@ class LearnerHandler : public OutputHandler<XGBoostModelHandle> {
 };
 
 /*! \brief handler for XGBoost checkpoint */
-class XGBoostCheckpointHandler : public OutputHandler<XGBoostModelHandle> {
+class XGBoostCheckpointHandler : public OutputHandler<ParsedXGBoostModel> {
  public:
-  using OutputHandler<XGBoostModelHandle>::OutputHandler;
+  using OutputHandler<ParsedXGBoostModel>::OutputHandler;
   bool StartArray() override;
   bool StartObject() override;
 };
 
 /*! \brief handler for XGBoostModel objects from XGBoost schema */
-class XGBoostModelHandler : public OutputHandler<XGBoostModelHandle> {
+class XGBoostModelHandler : public OutputHandler<ParsedXGBoostModel> {
  public:
-  using OutputHandler<XGBoostModelHandle>::OutputHandler;
+  using OutputHandler<ParsedXGBoostModel>::OutputHandler;
   bool StartArray() override;
   bool StartObject() override;
   bool EndObject(std::size_t memberCount) override;
 };
 
 /*! \brief handler for root object of XGBoost schema*/
-class RootHandler : public OutputHandler<std::unique_ptr<treelite::Model>> {
+class RootHandler : public OutputHandler<ParsedXGBoostModel> {
  public:
-  using OutputHandler<std::unique_ptr<treelite::Model>>::OutputHandler;
+  using OutputHandler<ParsedXGBoostModel>::OutputHandler;
   bool StartObject() override;
- private:
-  XGBoostModelHandle handle;
 };
 
 /*! \brief handler which delegates JSON parsing to stack of delegates*/
@@ -380,7 +380,7 @@ class DelegatedHandler
   void pop_delegate() override {
     delegates.pop();
   }
-  std::unique_ptr<treelite::Model> get_result();
+  ParsedXGBoostModel get_result();
   bool Null();
   bool Bool(bool b);
   bool Int(int i);
@@ -396,10 +396,15 @@ class DelegatedHandler
   bool EndArray(std::size_t elementCount);
 
  private:
-  DelegatedHandler() : delegates{}, result{treelite::Model::Create<float, float>()} {};
+  DelegatedHandler()
+  : delegates{},
+    result{treelite::Model::Create<float, float>(), nullptr, {}, {}, ""}
+  {
+    result.model = dynamic_cast<treelite::ModelImpl<float, float>*>(result.model_ptr.get());
+  }
 
   std::stack<std::shared_ptr<BaseHandler>> delegates;
-  std::unique_ptr<treelite::Model> result;
+  ParsedXGBoostModel result;
 };
 
 }  // namespace details

src/frontend/xgboost_json.cc

@@ -312,8 +312,8 @@ bool RegTreeHandler::EndObject(std::size_t) {
 bool GBTreeModelHandler::StartArray() {
   return (push_key_handler<ArrayHandler<treelite::Tree<float, float>, RegTreeHandler>,
                            std::vector<treelite::Tree<float, float>>>(
-                               "trees", output.trees) ||
-          push_key_handler<IgnoreHandler>("tree_info"));
+                               "trees", output.model->trees) ||
+          push_key_handler<ArrayHandler<int>, std::vector<int>>("tree_info", output.tree_info));
 }
 
 bool GBTreeModelHandler::StartObject() {
@@ -338,10 +338,9 @@ bool GradientBoosterHandler::String(const char *str,
   }
 }
 bool GradientBoosterHandler::StartObject() {
-  if (push_key_handler<GBTreeModelHandler, treelite::ModelImpl<float, float>>("model", output)) {
+  if (push_key_handler<GBTreeModelHandler, ParsedXGBoostModel>("model", output)) {
     return true;
-  } else if (push_key_handler<GradientBoosterHandler, treelite::ModelImpl<float, float>>("gbtree",
-                                                                                         output)) {
+  } else if (push_key_handler<GradientBoosterHandler, ParsedXGBoostModel>("gbtree", output)) {
     // "dart" booster contains a standard gbtree under ["gradient_booster"]["gbtree"]["model"].
     return true;
   } else {
@@ -356,11 +355,12 @@ bool GradientBoosterHandler::StartArray() {
 bool GradientBoosterHandler::EndObject(std::size_t memberCount) {
   if (name == "dart" && !weight_drop.empty()) {
     // Fold weight drop into leaf value for dart models.
-    TREELITE_CHECK_EQ(output.trees.size(), weight_drop.size());
-    for (size_t i = 0; i < output.trees.size(); ++i) {
-      for (int nid = 0; nid < output.trees[i].num_nodes; ++nid) {
-        if (output.trees[i].IsLeaf(nid)) {
-          output.trees[i].SetLeaf(nid, weight_drop[i] * output.trees[i].LeafValue(nid));
+    auto& trees = output.model->trees;
+    TREELITE_CHECK_EQ(trees.size(), weight_drop.size());
+    for (size_t i = 0; i < trees.size(); ++i) {
+      for (int nid = 0; nid < trees[i].num_nodes; ++nid) {
+        if (trees[i].IsLeaf(nid)) {
+          trees[i].SetLeaf(nid, weight_drop[i] * trees[i].LeafValue(nid));
         }
       }
     }
@@ -415,8 +415,8 @@ bool LearnerHandler::StartObject() {
   // "attributes" key is not documented in schema
   return (push_key_handler<LearnerParamHandler, treelite::ModelImpl<float, float>>(
               "learner_model_param", *output.model) ||
-          push_key_handler<GradientBoosterHandler, treelite::ModelImpl<float, float>>(
-              "gradient_booster", *output.model) ||
+          push_key_handler<GradientBoosterHandler, ParsedXGBoostModel>(
+              "gradient_booster", output) ||
           push_key_handler<ObjectiveHandler, std::string>("objective", objective) ||
           push_key_handler<IgnoreHandler>("attributes"));
 }
@@ -442,7 +442,7 @@ bool XGBoostCheckpointHandler::StartArray() {
 }
 
 bool XGBoostCheckpointHandler::StartObject() {
-  return push_key_handler<LearnerHandler, XGBoostModelHandle>("learner", output);
+  return push_key_handler<LearnerHandler, ParsedXGBoostModel>("learner", output);
 }
 
 /******************************************************************************
@@ -454,9 +454,9 @@ bool XGBoostModelHandler::StartArray() {
 }
 
 bool XGBoostModelHandler::StartObject() {
-  return (push_key_handler<LearnerHandler, XGBoostModelHandle>("learner", output) ||
+  return (push_key_handler<LearnerHandler, ParsedXGBoostModel>("learner", output) ||
           push_key_handler<IgnoreHandler>("Config") ||
-          push_key_handler<XGBoostCheckpointHandler, XGBoostModelHandle>("Model", output));
+          push_key_handler<XGBoostCheckpointHandler, ParsedXGBoostModel>("Model", output));
 }
 
 bool XGBoostModelHandler::EndObject(std::size_t memberCount) {
@@ -489,14 +489,13 @@ bool XGBoostModelHandler::EndObject(std::size_t memberCount) {
  * RootHandler
  * ***************************************************************************/
 bool RootHandler::StartObject() {
-  handle = {dynamic_cast<treelite::ModelImpl<float, float>*>(output.get()), {}, ""};
-  return push_handler<XGBoostModelHandler, XGBoostModelHandle>(handle);
+  return push_handler<XGBoostModelHandler, ParsedXGBoostModel>(output);
 }
 
 /******************************************************************************
  * DelegatedHandler
  * ***************************************************************************/
-std::unique_ptr<treelite::Model> DelegatedHandler::get_result() { return std::move(result); }
+ParsedXGBoostModel DelegatedHandler::get_result() { return std::move(result); }
 bool DelegatedHandler::Null() { return delegates.top()->Null(); }
 bool DelegatedHandler::Bool(bool b) { return delegates.top()->Bool(b); }
 bool DelegatedHandler::Int(int i) { return delegates.top()->Int(i); }
@@ -541,6 +540,37 @@ std::unique_ptr<treelite::Model> ParseStream(std::unique_ptr<StreamType> input_s
                         << "Parsing error at offset " << offset << ": "
                         << rapidjson::GetParseError_En(error_code) << "\n" << diagnostic;
   }
-  return handler->get_result();
+  treelite::details::ParsedXGBoostModel parsed = handler->get_result();
+
+  // Special handling for multi-class classifier with num_parallel_tree > 1
+  if (parsed.model->task_param.grove_per_class && parsed.model->task_param.num_class > 2) {
+    // Infer num_parallel_tree
+    unsigned num_parallel_tree = 0;
+    for (int e : parsed.tree_info) {
+      if (e != 0) {
+        break;
+      }
+      ++num_parallel_tree;
+    }
+    if (num_parallel_tree > 1) {
+      // Re-order trees to recover the grove-per-class layout.
+      // The prediction for the i-th class is determined by the trees whose index is congruent
+      // to [i] modulo [num_class].
+      // Currently, the trees' association with classes is as follows:
+      // 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2
+      // We need to re-order them as follows:
+      // 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2
+      std::vector<treelite::Tree<float, float>> new_trees;
+      std::size_t num_tree = parsed.model->trees.size();
+      for (std::size_t c = 0; c < num_parallel_tree; ++c) {
+        for (std::size_t tree_id = c; tree_id < num_tree; tree_id += num_parallel_tree) {
+          new_trees.push_back(std::move(parsed.model->trees[tree_id]));
+        }
+      }
+      TREELITE_CHECK_EQ(new_trees.size(), num_tree);
+      parsed.model->trees = std::move(new_trees);
+    }
+  }
+  return std::move(parsed.model_ptr);
 }
 }  // anonymous namespace

tests/cpp/test_frontend.cc

@@ -438,7 +438,9 @@ TEST(GBTreeModelHandlerSuite, GBTreeModelHandler) {
   std::shared_ptr<MockDelegator> delegator =
     std::make_shared<MockDelegator>();
 
-  ModelImpl<float, float> output;
+  details::ParsedXGBoostModel output{
+      Model::Create<float, float>(), nullptr, {}, {}, ""};
+  output.model = dynamic_cast<ModelImpl<float, float>*>(output.model_ptr.get());
   details::GBTreeModelHandler wrapped_handler {delegator, output};
   MockObjectStarter handler {delegator, wrapped_handler};
 
@@ -448,6 +450,33 @@ TEST(GBTreeModelHandlerSuite, GBTreeModelHandler) {
   reader.Parse(input_stream, handler);
 }
 
+TEST(GBTreeModelHandlerSuite, TreeInfoField) {
+  class GBTreeModelHandlerWrapper : public details::BaseHandler {
+   public:
+    using BaseHandler::BaseHandler;
+    bool StartObject() override {
+      push_handler<details::GBTreeModelHandler, details::ParsedXGBoostModel>(output);
+      return true;
+    }
+    details::ParsedXGBoostModel output;
+  };
+  auto handler = details::DelegatedHandler::create_empty();
+  auto wrapped_handler = std::make_shared<GBTreeModelHandlerWrapper>(handler);
+  wrapped_handler->output.model_ptr = Model::Create<float, float>();
+  wrapped_handler->output.model =
+      dynamic_cast<ModelImpl<float, float>*>(wrapped_handler->output.model_ptr.get());
+  handler->push_delegate(wrapped_handler);
+  rapidjson::Reader reader;
+
+  std::string json_str = "{\"trees\": [], \"tree_info\": [0, 1, 2]}";
+  auto input_stream = rapidjson::MemoryStream(json_str.c_str(),
+                                              json_str.size());
+  ASSERT_TRUE(reader.Parse(input_stream, *handler));
+
+  std::vector<int> expected_tree_info{0, 1, 2};
+  ASSERT_EQ(wrapped_handler->output.tree_info, expected_tree_info);
+}
+
 /******************************************************************************
  * GradientBoosterHandler
  * ***************************************************************************/
@@ -458,7 +487,9 @@ TEST(GradientBoosterHandlerSuite, GradientBoosterHandler) {
   std::shared_ptr<MockDelegator> delegator =
     std::make_shared<MockDelegator>();
 
-  ModelImpl<float, float> output;
+  details::ParsedXGBoostModel output{
+      Model::Create<float, float>(), nullptr, {}, {}, ""};
+  output.model = dynamic_cast<ModelImpl<float, float>*>(output.model_ptr.get());
   details::GradientBoosterHandler wrapped_handler {delegator, output};
   MockObjectStarter handler {delegator, wrapped_handler};
 
@@ -525,8 +556,9 @@ TEST(XGBoostModelHandlerSuite, XGBoostModelHandler) {
   std::shared_ptr<MockDelegator> delegator =
     std::make_shared<MockDelegator>();
 
-  ModelImpl<float, float> output_model;
-  details::XGBoostModelHandle output{&output_model, {}, ""};
+  details::ParsedXGBoostModel output{
+      Model::Create<float, float>(), nullptr, {}, {}, ""};
+  output.model = dynamic_cast<ModelImpl<float, float>*>(output.model_ptr.get());
   details::XGBoostModelHandler wrapped_handler{delegator, output};
   MockObjectStarter handler{delegator, wrapped_handler};
 
@@ -547,7 +579,7 @@ TEST(RootHandlerSuite, RootHandler) {
   std::shared_ptr<MockDelegator> delegator =
     std::make_shared<MockDelegator>();
 
-  std::unique_ptr<treelite::Model> output;
+  details::ParsedXGBoostModel output;
   details::RootHandler wrapped_handler {delegator, output};
   MockObjectStarter handler {delegator, wrapped_handler};