GPUTreeShap (#6038)

tests/cpp/predictor/test_gpu_predictor.cu

@@ -163,5 +163,61 @@ TEST(GPUPredictor, MGPU_InplacePredict) {  // NOLINT
 TEST(GpuPredictor, LesserFeatures) {
   TestPredictionWithLesserFeatures("gpu_predictor");
 }
+// Very basic test of empty model
+TEST(GPUPredictor, ShapStump) {
+  cudaSetDevice(0);
+  LearnerModelParam param;
+  param.num_feature = 1;
+  param.num_output_group = 1;
+  param.base_score = 0.5;
+  gbm::GBTreeModel model(&param);
+  std::vector<std::unique_ptr<RegTree>> trees;
+  trees.push_back(std::unique_ptr<RegTree>(new RegTree));
+  model.CommitModel(std::move(trees), 0);
+
+  auto gpu_lparam = CreateEmptyGenericParam(0);
+  std::unique_ptr<Predictor> gpu_predictor =
+      std::unique_ptr<Predictor>(Predictor::Create("gpu_predictor", &gpu_lparam));
+  gpu_predictor->Configure({});
+  std::vector<float > phis;
+    auto dmat = RandomDataGenerator(3, 1, 0).GenerateDMatrix();
+  gpu_predictor->PredictContribution(dmat.get(), &phis, model);
+  EXPECT_EQ(phis[0], 0.0);
+  EXPECT_EQ(phis[1], param.base_score);
+  EXPECT_EQ(phis[2], 0.0);
+  EXPECT_EQ(phis[3], param.base_score);
+  EXPECT_EQ(phis[4], 0.0);
+  EXPECT_EQ(phis[5], param.base_score);
+}
+TEST(GPUPredictor, Shap) {
+  LearnerModelParam param;
+  param.num_feature = 1;
+  param.num_output_group = 1;
+  param.base_score = 0.5;
+  gbm::GBTreeModel model(&param);
+  std::vector<std::unique_ptr<RegTree>> trees;
+  trees.push_back(std::unique_ptr<RegTree>(new RegTree));
+  trees[0]->ExpandNode(0, 0, 0.5, true, 1.0, -1.0, 1.0, 0.0, 5.0, 2.0, 3.0);
+  model.CommitModel(std::move(trees), 0);
+
+  auto gpu_lparam = CreateEmptyGenericParam(0);
+  auto cpu_lparam = CreateEmptyGenericParam(-1);
+  std::unique_ptr<Predictor> gpu_predictor =
+      std::unique_ptr<Predictor>(Predictor::Create("gpu_predictor", &gpu_lparam));
+  std::unique_ptr<Predictor> cpu_predictor =
+      std::unique_ptr<Predictor>(Predictor::Create("cpu_predictor", &cpu_lparam));
+  gpu_predictor->Configure({});
+  cpu_predictor->Configure({});
+  std::vector<float > phis;
+  std::vector<float > cpu_phis;
+  auto dmat = RandomDataGenerator(3, 1, 0).GenerateDMatrix();
+  gpu_predictor->PredictContribution(dmat.get(), &phis, model);
+  cpu_predictor->PredictContribution(dmat.get(), &cpu_phis, model);
+  for(auto i = 0ull; i < phis.size(); i++)
+  {
+    EXPECT_NEAR(cpu_phis[i], phis[i], 1e-3);
+  }
+}
+
 }  // namespace predictor
 }  // namespace xgboost

tests/python-gpu/test_gpu_prediction.py

@@ -4,6 +4,7 @@ import pytest
 
 import numpy as np
 import xgboost as xgb
+from hypothesis import given, strategies, assume, settings, note
 
 sys.path.append("tests/python")
 import testing as tm
@@ -11,6 +12,12 @@ from test_predict import run_threaded_predict  # noqa
 
 rng = np.random.RandomState(1994)
 
+shap_parameter_strategy = strategies.fixed_dictionaries({
+    'max_depth': strategies.integers(0, 11),
+    'max_leaves': strategies.integers(0, 256),
+    'num_parallel_tree': strategies.sampled_from([1, 10]),
+})
+
 
 class TestGPUPredict(unittest.TestCase):
     def test_predict(self):
@@ -149,7 +156,8 @@ class TestGPUPredict(unittest.TestCase):
 
         # Don't do this on Windows, see issue #5793
         if sys.platform.startswith("win"):
-            pytest.skip('Multi-threaded in-place prediction with cuPy is not working on Windows')
+            pytest.skip(
+                'Multi-threaded in-place prediction with cuPy is not working on Windows')
         for i in range(10):
             run_threaded_predict(X, rows, predict_dense)
 
@@ -185,3 +193,24 @@ class TestGPUPredict(unittest.TestCase):
 
         for i in range(10):
             run_threaded_predict(X, rows, predict_df)
+
+    @given(strategies.integers(1, 200),
+           tm.dataset_strategy, shap_parameter_strategy, strategies.booleans())
+    @settings(deadline=None)
+    def test_shap(self, num_rounds, dataset, param, all_rows):
+        param.update({"predictor": "gpu_predictor", "gpu_id": 0})
+        param = dataset.set_params(param)
+        dmat = dataset.get_dmat()
+        bst = xgb.train(param, dmat, num_rounds)
+        if all_rows:
+            test_dmat = xgb.DMatrix(dataset.X, dataset.y, dataset.w, dataset.margin)
+        else:
+            test_dmat = xgb.DMatrix(dataset.X[0:1, :])
+        shap = bst.predict(test_dmat, pred_contribs=True)
+        bst.set_param({"predictor": "cpu_predictor"})
+        cpu_shap = bst.predict(test_dmat, pred_contribs=True)
+        margin = bst.predict(test_dmat, output_margin=True)
+        assert np.allclose(shap, cpu_shap, 1e-3, 1e-3)
+        # feature contributions should add up to predictions
+        assume(len(dataset.y) > 0)
+        assert np.allclose(np.sum(shap, axis=len(shap.shape) - 1), margin, 1e-3, 1e-3)

tests/python/testing.py

@@ -131,6 +131,7 @@ class TestDataset:
         self.metric = metric
         self.X, self.y = get_dataset()
         self.w = None
+        self.margin = None
 
     def set_params(self, params_in):
         params_in['objective'] = self.objective
@@ -140,13 +141,13 @@ class TestDataset:
         return params_in
 
     def get_dmat(self):
-        return xgb.DMatrix(self.X, self.y, self.w)
+        return xgb.DMatrix(self.X, self.y, self.w, base_margin=self.margin)
 
     def get_device_dmat(self):
         w = None if self.w is None else cp.array(self.w)
         X = cp.array(self.X, dtype=np.float32)
         y = cp.array(self.y, dtype=np.float32)
-        return xgb.DeviceQuantileDMatrix(X, y, w)
+        return xgb.DeviceQuantileDMatrix(X, y, w, base_margin=self.margin)
 
     def get_external_dmat(self):
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -157,7 +158,7 @@ class TestDataset:
             uri = path + '?format=csv&label_column=0#tmptmp_'
             # The uri looks like:
             # 'tmptmp_1234.csv?format=csv&label_column=0#tmptmp_'
-            return xgb.DMatrix(uri, weight=self.w)
+            return xgb.DMatrix(uri, weight=self.w, base_margin=self.margin)
 
     def __repr__(self):
         return self.name
@@ -206,16 +207,23 @@ _unweighted_datasets_strategy = strategies.sampled_from(
 
 
 @strategies.composite
-def _dataset_and_weight(draw):
+def _dataset_weight_margin(draw):
     data = draw(_unweighted_datasets_strategy)
     if draw(strategies.booleans()):
         data.w = draw(arrays(np.float64, (len(data.y)), elements=strategies.floats(0.1, 2.0)))
+    if draw(strategies.booleans()):
+        num_class = 1
+        if data.objective == "multi:softmax":
+            num_class = int(np.max(data.y) + 1)
+        data.margin = draw(
+            arrays(np.float64, (len(data.y) * num_class), elements=strategies.floats(0.5, 1.0)))
+
     return data
 
 
 # A strategy for drawing from a set of example datasets
 # May add random weights to the dataset
-dataset_strategy = _dataset_and_weight()
+dataset_strategy = _dataset_weight_margin()
 
 
 def non_increasing(L, tolerance=1e-4):