Files
2026-07-13 13:22:52 +08:00

386 lines
13 KiB
Python

"""Test gpu accelerated tree functions."""
import numpy as np
import pandas as pd
import pytest
import sklearn
import shap
from shap.explainers._tree import SingleTree, TreeEnsemble
from shap.utils import assert_import
try:
assert_import("cext_gpu")
except ImportError:
pytestmark = pytest.mark.skip("cuda module not built")
def test_front_page_xgboost():
xgboost = pytest.importorskip("xgboost")
# load JS visualization code to notebook
shap.initjs()
# train XGBoost model
X, y = shap.datasets.california(n_points=500)
model = xgboost.train({"learning_rate": 0.01}, xgboost.DMatrix(X, label=y), 100)
# explain the model's predictions using SHAP values
explainer = shap.GPUTreeExplainer(model)
shap_values = explainer.shap_values(X)
# visualize the first prediction's explanation
shap.force_plot(explainer.expected_value, shap_values[0, :], X.iloc[0, :])
# visualize the training set predictions
shap.force_plot(explainer.expected_value, shap_values, X)
# create a SHAP dependence plot to show the effect of a single feature across the whole dataset
shap.dependence_plot(5, shap_values, X, show=False)
shap.dependence_plot("Longitude", shap_values, X, show=False)
# summarize the effects of all the features
shap.summary_plot(shap_values, X, show=False)
rs = np.random.RandomState(15921)
n = 100
m = 4
datasets = {
"regression": (rs.randn(n, m), rs.randn(n)),
"binary": (rs.randn(n, m), rs.binomial(1, 0.5, n)),
"multiclass": (rs.randn(n, m), rs.randint(0, 5, n)),
}
def task_xfail(func):
def inner():
return pytest.param(func(), marks=pytest.mark.xfail)
return inner
def xgboost_base():
try:
import xgboost
except ImportError:
return pytest.param("xgboost.XGBRegressor", marks=pytest.mark.skip)
X, y = datasets["regression"]
model = xgboost.XGBRegressor(tree_method="hist")
model.fit(X, y)
return model.get_booster(), X, model.predict(X)
def xgboost_regressor():
try:
import xgboost
except ImportError:
return pytest.param("xgboost.XGBRegressor", marks=pytest.mark.skip)
X, y = datasets["regression"]
model = xgboost.XGBRegressor()
model.fit(X, y)
return model, X, model.predict(X)
def xgboost_binary_classifier():
try:
import xgboost
except ImportError:
return pytest.param("xgboost.XGBClassifier", marks=pytest.mark.skip)
X, y = datasets["binary"]
model = xgboost.XGBClassifier(eval_metric="error")
model.fit(X, y)
return model, X, model.predict(X, output_margin=True)
def xgboost_multiclass_classifier():
try:
import xgboost
except ImportError:
return pytest.param("xgboost.XGBClassifier", marks=pytest.mark.skip)
X, y = datasets["multiclass"]
model = xgboost.XGBClassifier()
model.fit(X, y)
return model, X, model.predict(X, output_margin=True)
def test_xgboost_cat_unsupported() -> None:
xgboost = pytest.importorskip("xgboost")
X, y = shap.datasets.adult()
X["Workclass"] = X["Workclass"].astype("category")
clf = xgboost.XGBClassifier(n_estimators=2, enable_categorical=True, device="cuda")
clf.fit(X, y)
# Tests for both CPU and GPU in one place
# Prefer an explict error over silent invalid values.
gpu_ex = shap.GPUTreeExplainer(clf, X, feature_perturbation="interventional")
with pytest.raises(NotImplementedError, match="Categorical"):
gpu_ex.shap_values(X)
ex = shap.TreeExplainer(clf, X, feature_perturbation="interventional")
with pytest.raises(NotImplementedError, match="Categorical"):
ex.shap_values(X)
def lightgbm_base():
try:
import lightgbm
except ImportError:
return pytest.param("lightgbm.LGBMRegressor", marks=pytest.mark.skip)
X, y = datasets["regression"]
model = lightgbm.LGBMRegressor(n_jobs=1)
model.fit(X, y)
return model.booster_, X, model.predict(X)
def lightgbm_regression():
try:
import lightgbm
except ImportError:
return pytest.param("lightgbm.LGBMRegressor", marks=pytest.mark.skip)
X, y = datasets["regression"]
model = lightgbm.LGBMRegressor(n_jobs=1)
model.fit(X, y)
return model, X, model.predict(X)
def lightgbm_binary_classifier():
try:
import lightgbm
except ImportError:
return pytest.param("lightgbm.LGBMClassifier", marks=pytest.mark.skip)
X, y = datasets["binary"]
model = lightgbm.LGBMClassifier(n_jobs=1)
model.fit(X, y)
return model, X, model.predict(X, raw_score=True)
def lightgbm_multiclass_classifier():
try:
import lightgbm
except ImportError:
return pytest.param("lightgbm.LGBMClassifier", marks=pytest.mark.skip)
X, y = datasets["multiclass"]
model = lightgbm.LGBMClassifier(n_jobs=1)
model.fit(X, y)
return model, X, model.predict(X, raw_score=True)
def rf_regressor():
X, y = datasets["regression"]
model = sklearn.ensemble.RandomForestRegressor()
model.fit(X, y)
return model, X, model.predict(X)
def rf_binary_classifier():
X, y = datasets["binary"]
model = sklearn.ensemble.RandomForestClassifier()
model.fit(X, y)
return model, X, model.predict_proba(X)
def rf_multiclass_classifier():
X, y = datasets["multiclass"]
model = sklearn.ensemble.RandomForestClassifier()
model.fit(X, y)
return model, X, model.predict_proba(X)
tasks = [
xgboost_base(),
xgboost_regressor(),
xgboost_binary_classifier(),
xgboost_multiclass_classifier(),
lightgbm_base(),
lightgbm_regression(),
lightgbm_binary_classifier(),
lightgbm_multiclass_classifier(),
rf_binary_classifier(),
rf_regressor(),
rf_multiclass_classifier(),
]
# pretty print tasks
def idfn(task):
if isinstance(task, str):
return task
model, _, _ = task
return type(model).__module__ + "." + type(model).__qualname__
def assert_gpu_matches_cpu(task, feature_perturbation, X=None):
model, background, _ = task
if X is None:
X = background
gpu_ex = shap.GPUTreeExplainer(model, background, feature_perturbation=feature_perturbation)
ex = shap.TreeExplainer(model, background, feature_perturbation=feature_perturbation)
host_shap = ex.shap_values(X, check_additivity=True)
gpu_shap = gpu_ex.shap_values(X, check_additivity=True)
# todo: this should actually happen in the GPUTreeExplainer
if np.array(gpu_shap).ndim == 3:
gpu_shap = np.moveaxis(np.array(gpu_shap), [0, 1, 2], [2, 0, 1])
else:
gpu_shap = np.array(gpu_shap, copy=False)
# Check outputs roughly the same as CPU algorithm
assert np.allclose(ex.expected_value, gpu_ex.expected_value, 1e-3, 1e-3)
assert np.allclose(host_shap, gpu_shap, 1e-3, 1e-3)
@pytest.mark.parametrize("task", tasks, ids=idfn)
@pytest.mark.parametrize("feature_perturbation", ["interventional", "tree_path_dependent"])
def test_gpu_tree_explainer_shap(task, feature_perturbation):
assert_gpu_matches_cpu(task, feature_perturbation)
def test_gpu_tree_explainer_shap_with_missing_values():
task = xgboost_base()
X = task[1].copy()
rows = np.arange(0, X.shape[0], 10)
X[rows, 0] = np.nan
for feature_perturbation in ["interventional", "tree_path_dependent"]:
assert_gpu_matches_cpu(task, feature_perturbation, X)
@pytest.mark.parametrize("task", tasks, ids=idfn)
@pytest.mark.parametrize("feature_perturbation", ["tree_path_dependent"])
def test_gpu_tree_explainer_shap_interactions(task, feature_perturbation):
model, X, margin = task
ex = shap.GPUTreeExplainer(model, X, feature_perturbation=feature_perturbation)
shap_values = np.array(ex.shap_interaction_values(X), copy=False)
assert np.allclose(np.sum(shap_values, axis=(1, 2)) + ex.expected_value, margin, atol=1e-4)
@pytest.mark.parametrize("use_interactions", [False, True])
def test_lightgbm_categorical_split(use_interactions):
# GH 480
"""Checks that shap interaction values are computed without error when the LightGBM model has categorical splits."""
lightgbm = pytest.importorskip("lightgbm")
X, y = shap.datasets.california(n_points=10000)
# Add HouseAgeGroup categorical variable
target_variable = "HouseAge"
X["HouseAgeGroup"] = pd.cut(
X[target_variable],
bins=[-float("inf"), 17, 27, 37, float("inf")],
labels=[0, 1, 2, 3],
right=False,
).astype(int)
model = lightgbm.LGBMRegressor(n_estimators=400, max_cat_to_onehot=1)
model.fit(
X, y, categorical_feature=[X.columns.get_loc("HouseAgeGroup")]
) # Set HouseAgeGroup as categorical variable
preds = model.predict(X, raw_score=True)
explainer = shap.GPUTreeExplainer(model)
if use_interactions:
# Check SHAP interaction values sum to model output
shap_interaction_values = explainer.shap_interaction_values(X.iloc[:10, :])
assert np.allclose(shap_interaction_values.sum(axis=(1, 2)) + explainer.expected_value, preds[:10], atol=1e-4)
else:
shap_values = explainer.shap_values(X.iloc[:10, :])
assert np.allclose(shap_values.sum(axis=1) + explainer.expected_value, preds[:10], atol=1e-4)
def test_categorical_split_cpu_gpu_equivalence():
"""
Check consistency with a dummy tree that a single categorical split yields the same results on GPU and CPU.
"""
tree = {
"children_left": np.array([1, -1, -1], dtype=np.int32),
"children_right": np.array([2, -1, -1], dtype=np.int32),
"children_default": np.array([1, -1, -1], dtype=np.int32),
"features": np.array([0, -1, -1], dtype=np.int32),
"thresholds": np.array([2.0, 0.0, 0.0], dtype=np.float64),
"values": np.array([[0.8], [2.0], [-1.0]], dtype=np.float64),
"node_sample_weight": np.array([100.0, 60.0, 40.0], dtype=np.float64),
}
single_tree = SingleTree(tree)
single_tree.threshold_types = np.array([1, 0, 0], dtype=np.int32)
ensemble = TreeEnsemble([single_tree], model_output="raw")
ensemble.tree_output = "raw_value"
ensemble.objective = "squared_error"
X = np.array([[0.0], [1.0], [2.0], [3.0]])
cpu_explainer = shap.TreeExplainer(ensemble, feature_perturbation="tree_path_dependent")
gpu_explainer = shap.GPUTreeExplainer(ensemble)
shap_values_cpu = cpu_explainer.shap_values(X, check_additivity=False)
shap_values_gpu = gpu_explainer.shap_values(X, check_additivity=False)
np.testing.assert_allclose(shap_values_gpu, shap_values_cpu, atol=1e-5)
def test_categorical_split_matches_binary_feature():
"""
Tests that using the categorical feature for SHAP value computation gives the same result as using a binary feature that routes the same way. We compare values computed on gpu and cpu here to check consistency.
"""
children_left = np.array([1, -1, -1], dtype=np.int32)
children_right = np.array([2, -1, -1], dtype=np.int32)
children_default = np.array([1, -1, -1], dtype=np.int32)
features = np.array([0, -1, -1], dtype=np.int32)
values = np.array([[0.8], [2.0], [-1.0]], dtype=np.float64)
node_sample_weight = np.array([100.0, 60.0, 40.0], dtype=np.float64)
cat_tree = {
"children_left": children_left,
"children_right": children_right,
"children_default": children_default,
"features": features,
"thresholds": np.array([1.0, 0.0, 0.0], dtype=np.float64),
"values": values,
"node_sample_weight": node_sample_weight,
}
cat_single = SingleTree(cat_tree)
cat_single.threshold_types = np.array([1, 0, 0], dtype=np.int32)
cat_ensemble = TreeEnsemble([cat_single], model_output="raw")
cat_ensemble.tree_output = "raw_value"
cat_ensemble.objective = "squared_error"
bin_tree = {
"children_left": children_left,
"children_right": children_right,
"children_default": children_default,
"features": features,
"thresholds": np.array([0.5, 0.0, 0.0], dtype=np.float64),
"values": values,
"node_sample_weight": node_sample_weight,
}
bin_single = SingleTree(bin_tree)
bin_ensemble = TreeEnsemble([bin_single], model_output="raw")
bin_ensemble.tree_output = "raw_value"
bin_ensemble.objective = "squared_error"
X_cat = np.array([[1.0], [2.0], [1.0], [2.0]])
X_bin = X_cat - 1.0
cat_cpu = shap.TreeExplainer(cat_ensemble, feature_perturbation="tree_path_dependent").shap_values(
X_cat, check_additivity=False
)
cat_gpu = shap.GPUTreeExplainer(cat_ensemble).shap_values(X_cat, check_additivity=False)
np.testing.assert_allclose(cat_gpu, cat_cpu, atol=1e-5)
bin_cpu = shap.TreeExplainer(bin_ensemble, feature_perturbation="tree_path_dependent").shap_values(
X_bin, check_additivity=False
)
bin_gpu = shap.GPUTreeExplainer(bin_ensemble).shap_values(X_bin, check_additivity=False)
np.testing.assert_allclose(bin_gpu, bin_cpu, atol=1e-5)
np.testing.assert_allclose(cat_gpu, bin_gpu, atol=1e-5)
np.testing.assert_allclose(cat_cpu, cat_gpu, atol=1e-5)