Unify DART SHAP forwarding

Author: RAMitchell

src/gbm/gbtree.cc

@@ -797,23 +797,9 @@ class Dart : public GBTree {
     }
   }
 
-  void PredictContribution(DMatrix* p_fmat, HostDeviceVector<bst_float>* out_contribs,
-                           bst_layer_t layer_begin, bst_layer_t layer_end,
-                           bool approximate) override {
-    auto [tree_begin, tree_end] = detail::LayerToTree(model_, layer_begin, layer_end);
-    cpu_predictor_->PredictContribution(p_fmat, out_contribs, model_, tree_end, &weight_drop_,
-                                        approximate);
-  }
-
-  void PredictInteractionContributions(DMatrix* p_fmat, HostDeviceVector<float>* out_contribs,
-                                       bst_layer_t layer_begin, bst_layer_t layer_end,
-                                       bool approximate) override {
-    auto [tree_begin, tree_end] = detail::LayerToTree(model_, layer_begin, layer_end);
-    cpu_predictor_->PredictInteractionContributions(p_fmat, out_contribs, model_, tree_end,
-                                                    &weight_drop_, approximate);
-  }
-
  protected:
+  [[nodiscard]] std::vector<float> const* TreeWeights() const override { return &weight_drop_; }
+
   // commit new trees all at once
   void CommitModel(TreesOneIter&& new_trees) override {
     auto n_new_trees = model_.CommitModel(std::forward<TreesOneIter>(new_trees));

src/gbm/gbtree.h

@@ -303,8 +303,8 @@ class GBTree : public GradientBooster {
     auto [tree_begin, tree_end] = detail::LayerToTree(model_, layer_begin, layer_end);
     CHECK_EQ(tree_begin, 0) << "Predict contribution supports only iteration end: [0, "
                                "n_iteration), using model slicing instead.";
-    this->GetPredictor(false)->PredictContribution(p_fmat, out_contribs, model_, tree_end, nullptr,
-                                                   approximate);
+    this->GetPredictor(false)->PredictContribution(p_fmat, out_contribs, model_, tree_end,
+                                                   this->TreeWeights(), approximate);
   }
 
   void PredictInteractionContributions(DMatrix* p_fmat, HostDeviceVector<float>* out_contribs,
@@ -313,8 +313,8 @@ class GBTree : public GradientBooster {
     auto [tree_begin, tree_end] = detail::LayerToTree(model_, layer_begin, layer_end);
     CHECK_EQ(tree_begin, 0) << "Predict interaction contribution supports only iteration end: [0, "
                                "n_iteration), using model slicing instead.";
-    this->GetPredictor(false)->PredictInteractionContributions(p_fmat, out_contribs, model_,
-                                                               tree_end, nullptr, approximate);
+    this->GetPredictor(false)->PredictInteractionContributions(
+        p_fmat, out_contribs, model_, tree_end, this->TreeWeights(), approximate);
   }
 
   [[nodiscard]] std::vector<std::string> DumpModel(const FeatureMap& fmap, bool with_stats,
@@ -323,6 +323,8 @@ class GBTree : public GradientBooster {
   }
 
  protected:
+  [[nodiscard]] virtual std::vector<float> const* TreeWeights() const { return nullptr; }
+
   void BoostNewTrees(GradientContainer* gpair, DMatrix* p_fmat, int bst_group,
                      std::vector<HostDeviceVector<bst_node_t>>* out_position,
                      std::vector<std::unique_ptr<RegTree>>* ret);

tests/cpp/predictor/test_shap.cc

@@ -215,6 +215,38 @@ void CheckShapOutput(DMatrix* dmat, Args const& model_args) {
   CheckShapAdditivity(kRows, kCols, shap_interactions, margin_predt);
 }
 
+void CheckDartShapOutput(Context const* ctx) {
+  size_t constexpr kRows = 64, kCols = 8;
+  auto dmat = RandomDataGenerator(kRows, kCols, 0.0).Device(ctx->Device()).GenerateDMatrix();
+  SetLabels(dmat.get(), 1);
+
+  std::unique_ptr<Learner> learner{Learner::Create({dmat})};
+  learner->SetParams(Args{{"booster", "dart"},
+                          {"objective", "binary:logistic"},
+                          {"max_depth", "3"},
+                          {"rate_drop", "0.5"},
+                          {"sample_type", "uniform"},
+                          {"normalize_type", "tree"},
+                          {"device", ctx->IsSycl() ? "cpu" : ctx->DeviceName()}});
+  learner->Configure();
+  for (size_t i = 0; i < 4; ++i) {
+    learner->UpdateOneIter(i, dmat);
+  }
+
+  HostDeviceVector<float> margin_predt;
+  learner->Predict(dmat, true, &margin_predt, 0, 0, false, false, false, false, false);
+
+  HostDeviceVector<float> shap_values;
+  learner->Predict(dmat, false, &shap_values, 0, 0, false, false, true, false, false);
+  ASSERT_EQ(shap_values.Size(), kRows * (kCols + 1));
+  CheckShapAdditivity(kRows, kCols, shap_values, margin_predt);
+
+  HostDeviceVector<float> shap_interactions;
+  learner->Predict(dmat, false, &shap_interactions, 0, 0, false, false, false, false, true);
+  ASSERT_EQ(shap_interactions.Size(), kRows * (kCols + 1) * (kCols + 1));
+  CheckShapAdditivity(kRows, kCols, shap_interactions, margin_predt);
+}
+
 void CheckShapAdditivity(size_t rows, size_t cols, HostDeviceVector<float> const& shap_values,
                          HostDeviceVector<float> const& margin_predt) {
   auto const& h_shap = shap_values.ConstHostVector();
@@ -254,6 +286,11 @@ TEST(Predictor, ShapOutputCasesCPU) {
   }
 }
 
+TEST(Predictor, DartShapOutputCPU) {
+  Context ctx;
+  CheckDartShapOutput(&ctx);
+}
+
 TEST(Predictor, ApproxContribsBasic) {
   Context ctx;
   size_t constexpr kRows = 64;

tests/cpp/predictor/test_shap.cu

@@ -84,4 +84,9 @@ TEST(GPUPredictor, ShapOutputCasesGPU) {
     CheckShapOutput(dmat.get(), args);
   }
 }
+
+TEST(GPUPredictor, DartShapOutputGPU) {
+  auto ctx = MakeCUDACtx(0);
+  CheckDartShapOutput(&ctx);
+}
 }  // namespace xgboost

tests/cpp/predictor/test_shap.h

@@ -19,6 +19,7 @@ class Learner;
 
 namespace xgboost {
 void CheckShapOutput(DMatrix* dmat, Args const& model_args);
+void CheckDartShapOutput(Context const* ctx);
 void CheckShapAdditivity(size_t rows, size_t cols, HostDeviceVector<float> const& shap_values,
                          HostDeviceVector<float> const& margin_predt);