[EM] Support SHAP contribution with QDM. (#10724)

- Add GPU support.
- Add external memory support.
- Update the GPU tree shap.
This commit is contained in:
Jiaming Yuan
2024-08-22 05:25:10 +08:00
committed by GitHub
parent cb54374550
commit 142bdc73ec
13 changed files with 274 additions and 159 deletions

View File

@@ -343,32 +343,45 @@ class TestGPUPredict:
strategies.integers(1, 10), tm.make_dataset_strategy(), shap_parameter_strategy
)
@settings(deadline=None, max_examples=20, print_blob=True)
def test_shap(self, num_rounds, dataset, param):
def test_shap(self, num_rounds: int, dataset: tm.TestDataset, param: dict) -> None:
if dataset.name.endswith("-l1"): # not supported by the exact tree method
return
param.update({"tree_method": "hist", "device": "gpu:0"})
param = dataset.set_params(param)
dmat = dataset.get_dmat()
bst = xgb.train(param, dmat, num_rounds)
test_dmat = xgb.DMatrix(dataset.X, dataset.y, dataset.w, dataset.margin)
test_dmat = xgb.DMatrix(
dataset.X, dataset.y, weight=dataset.w, base_margin=dataset.margin
)
bst.set_param({"device": "gpu:0"})
shap = bst.predict(test_dmat, pred_contribs=True)
margin = bst.predict(test_dmat, output_margin=True)
assume(len(dataset.y) > 0)
assert np.allclose(np.sum(shap, axis=len(shap.shape) - 1), margin, 1e-3, 1e-3)
dmat = dataset.get_external_dmat()
shap = bst.predict(dmat, pred_contribs=True)
margin = bst.predict(dmat, output_margin=True)
assume(len(dataset.y) > 0)
assert np.allclose(np.sum(shap, axis=len(shap.shape) - 1), margin, 1e-3, 1e-3)
@given(
strategies.integers(1, 10), tm.make_dataset_strategy(), shap_parameter_strategy
)
@settings(deadline=None, max_examples=10, print_blob=True)
def test_shap_interactions(self, num_rounds, dataset, param):
def test_shap_interactions(
self, num_rounds: int, dataset: tm.TestDataset, param: dict
) -> None:
if dataset.name.endswith("-l1"): # not supported by the exact tree method
return
param.update({"tree_method": "hist", "device": "cuda:0"})
param = dataset.set_params(param)
dmat = dataset.get_dmat()
bst = xgb.train(param, dmat, num_rounds)
test_dmat = xgb.DMatrix(dataset.X, dataset.y, dataset.w, dataset.margin)
test_dmat = xgb.DMatrix(
dataset.X, dataset.y, weight=dataset.w, base_margin=dataset.margin
)
bst.set_param({"device": "cuda:0"})
shap = bst.predict(test_dmat, pred_interactions=True)
margin = bst.predict(test_dmat, output_margin=True)
@@ -380,6 +393,17 @@ class TestGPUPredict:
1e-3,
)
test_dmat = dataset.get_external_dmat()
shap = bst.predict(test_dmat, pred_interactions=True)
margin = bst.predict(test_dmat, output_margin=True)
assume(len(dataset.y) > 0)
assert np.allclose(
np.sum(shap, axis=(len(shap.shape) - 1, len(shap.shape) - 2)),
margin,
1e-3,
1e-3,
)
def test_shap_categorical(self):
X, y = tm.make_categorical(100, 20, 7, False)
Xy = xgb.DMatrix(X, y, enable_categorical=True)