fromurllib.errorimportHTTPErrorimportshapfromsklearn.datasetsimportfetch_california_housing,make_regressionimportxgboostasxgb# Fetch dataset using sklearntry:_data=fetch_california_housing(return_X_y=True)X=_data.datay=_data.targetfeature_names=_data.feature_namesprint(_data.DESCR)exceptHTTPError:# Use a synthetic dataset instead if we couldn'tX,y=make_regression(n_samples=20640,n_features=8,random_state=1234)feature_names=[f"f{i}"foriinrange(8)]num_round=500param={"eta":0.05,"max_depth":10,"tree_method":"hist","device":"cuda",}# GPU accelerated trainingdtrain=xgb.DMatrix(X,label=y,feature_names=feature_names)model=xgb.train(param,dtrain,num_round)# Compute shap values using GPU with xgboostmodel.set_param({"device":"cuda"})shap_values=model.predict(dtrain,pred_contribs=True)# Compute shap interaction values using GPUshap_interaction_values=model.predict(dtrain,pred_interactions=True)# shap will call the GPU accelerated version as long as the device parameter is set to# "cuda"explainer=shap.TreeExplainer(model)shap_values=explainer.shap_values(X)# visualize the first prediction's explanationshap.force_plot(explainer.expected_value,shap_values[0,:],X[0,:],feature_names=feature_names,matplotlib=True,)# Show a summary of feature importanceshap.summary_plot(shap_values,X,plot_type="bar",feature_names=feature_names)

DownloadJupyternotebook:gpu_tree_shap.ipynb

DownloadPythonsourcecode:gpu_tree_shap.py

Downloadzipped:gpu_tree_shap.zip