feature: Add support for Partial Dependence Plots(PDP) in SageMaker Clarify

Author: keerthanvasist

src/sagemaker/clarify.py

@@ -20,6 +20,8 @@
 import os
 import tempfile
 import re
+from typing import List
+
 from sagemaker.processing import ProcessingInput, ProcessingOutput, Processor
 from sagemaker import image_uris, s3, utils
 
@@ -292,7 +294,18 @@ class ExplainabilityConfig(ABC):
     @abstractmethod
     def get_explainability_config(self):
         """Returns config."""
-        return None
+
+
+class PDPConfig(ExplainabilityConfig):
+    def __init__(self, features=None, grid_resolution=None):
+        self.pdp_config = {}
+        if features is not None:
+            self.pdp_config["features"] = features
+        if grid_resolution is not None:
+            self.pdp_config["grid_resolution"] = grid_resolution
+
+    def get_explainability_config(self):
+        return {"pdp": copy.deepcopy(self.pdp_config)}
 
 
 class SHAPConfig(ExplainabilityConfig):
@@ -771,8 +784,9 @@ def run_explainability(
             data_config (:class:`~sagemaker.clarify.DataConfig`): Config of the input/output data.
             model_config (:class:`~sagemaker.clarify.ModelConfig`): Config of the model and its
                 endpoint to be created.
-            explainability_config (:class:`~sagemaker.clarify.ExplainabilityConfig`): Config of the
-                specific explainability method. Currently, only SHAP is supported.
+            explainability_config (:class:`~sagemaker.clarify.ExplainabilityConfig`| list of
+                :class:`~sagemaker.clarify.ExplainabilityConfig`): Config of the specific explainability method or a
+                 list of ExplainabilityConfig objects. Currently, SHAP and PDP are the two methods supported.
             model_scores(str|int|ModelPredictedLabelConfig):  Index or JSONPath location in the
                 model output for the predicted scores to be explained. This is not required if the
                 model output is a single score. Alternatively, an instance of
@@ -806,7 +820,21 @@ def run_explainability(
             predictor_config.update(predicted_label_config)
         else:
             _set(model_scores, "label", predictor_config)
-        analysis_config["methods"] = explainability_config.get_explainability_config()
+
+        explainability_methods = {}
+        if isinstance(explainability_config, List):
+            for config in explainability_config:
+                if not isinstance(config, ExplainabilityConfig):
+                    raise ValueError(
+                        f"Invalid input: Excepted ExplainabilityConfig, got {type(config)} instead"
+                    )
+                explain_config = config.get_explainability_config()
+                explainability_methods[list(explain_config.keys())[0]] = explain_config[
+                    list(explain_config.keys())[0]
+                ]
+        elif isinstance(explainability_config, ExplainabilityConfig):
+            explainability_methods = explainability_config.get_explainability_config()
+        analysis_config["methods"] = explainability_methods
         analysis_config["predictor"] = predictor_config
         if job_name is None:
             if self.job_name_prefix:

tests/integ/test_clarify.py

@@ -13,7 +13,6 @@
 
 from __future__ import print_function, absolute_import
 
-
 import json
 import math
 import numpy as np
@@ -31,14 +30,14 @@
     ModelConfig,
     ModelPredictedLabelConfig,
     SHAPConfig,
+    PDPConfig,
 )
 
 from sagemaker.amazon.linear_learner import LinearLearner, LinearLearnerPredictor
 from sagemaker import utils
 from tests import integ
 from tests.integ import timeout
 
-
 CLARIFY_DEFAULT_TIMEOUT_MINUTES = 15
 
 
@@ -177,6 +176,11 @@ def shap_config():
     )
 
 
+@pytest.fixture(scope="module")
+def pdp_config():
+    return PDPConfig(features=["F1"], grid_resolution=10)
+
+
 def test_pre_training_bias(clarify_processor, data_config, data_bias_config, sagemaker_session):
     with timeout.timeout(minutes=CLARIFY_DEFAULT_TIMEOUT_MINUTES):
         clarify_processor.run_pre_training_bias(
@@ -258,6 +262,57 @@ def test_shap(clarify_processor, data_config, model_config, shap_config, sagemak
         check_analysis_config(data_config, sagemaker_session, "shap")
 
 
+def test_pdp(clarify_processor, data_config, model_config, pdp_config, sagemaker_session):
+    with timeout.timeout(minutes=CLARIFY_DEFAULT_TIMEOUT_MINUTES):
+        clarify_processor.run_explainability(
+            data_config,
+            model_config,
+            pdp_config,
+            model_scores="score",
+            job_name=utils.unique_name_from_base("clarify-explainability-pdp"),
+            wait=True,
+        )
+        analysis_result_json = s3.S3Downloader.read_file(
+            data_config.s3_output_path + "/analysis.json",
+            sagemaker_session,
+        )
+        analysis_result = json.loads(analysis_result_json)
+        print(analysis_result)
+        assert analysis_result["explanations"]["pdp"][0]["feature_name"] == "F1"
+
+        check_analysis_config(data_config, sagemaker_session, "pdp")
+
+
+def test_shap_and_pdp(
+    clarify_processor, data_config, model_config, shap_config, pdp_config, sagemaker_session
+):
+    with timeout.timeout(minutes=CLARIFY_DEFAULT_TIMEOUT_MINUTES):
+        clarify_processor.run_explainability(
+            data_config,
+            model_config,
+            [shap_config, pdp_config],
+            model_scores="score",
+            job_name=utils.unique_name_from_base("clarify-explainability"),
+            wait=True,
+        )
+        analysis_result_json = s3.S3Downloader.read_file(
+            data_config.s3_output_path + "/analysis.json",
+            sagemaker_session,
+        )
+        analysis_result = json.loads(analysis_result_json)
+        print(analysis_result)
+        assert (
+            math.fabs(
+                analysis_result["explanations"]["kernel_shap"]["label0"]["global_shap_values"]["F2"]
+            )
+            <= 1
+        )
+        assert analysis_result["explanations"]["pdp"][0]["feature_name"] == "F1"
+
+        check_analysis_config(data_config, sagemaker_session, "pdp")
+        check_analysis_config(data_config, sagemaker_session, "shap")
+
+
 def check_analysis_config(data_config, sagemaker_session, method):
     analysis_config_json = s3.S3Downloader.read_file(
         data_config.s3_output_path + "/analysis_config.json",

tests/integ/test_clarify_model_monitor.py

@@ -42,6 +42,7 @@
     ModelConfig,
     ModelPredictedLabelConfig,
     SHAPConfig,
+    PDPConfig,
 )
 from sagemaker.model import Model
 

tests/unit/test_clarify.py

@@ -23,6 +23,7 @@
     ModelConfig,
     ModelPredictedLabelConfig,
     SHAPConfig,
+    PDPConfig,
 )
 from sagemaker import image_uris
 
@@ -268,6 +269,12 @@ def test_shap_config():
     assert expected_config == shap_config.get_explainability_config()
 
 
+def test_pdp_config():
+    pdp_config = PDPConfig(features=["f1", "f2"], grid_resolution=20)
+    expected_config = {"pdp": {"features": ["f1", "f2"], "grid_resolution": 20}}
+    assert expected_config == pdp_config.get_explainability_config()
+
+
 def test_invalid_shap_config():
     with pytest.raises(ValueError) as error:
         SHAPConfig(
@@ -374,6 +381,11 @@ def shap_config():
     )
 
 
+@pytest.fixture(scope="module")
+def pdp_config():
+    return PDPConfig(features=["F1", "F2"], grid_resolution=20)
+
+
 @patch("sagemaker.utils.name_from_base", return_value=JOB_NAME)
 def test_pre_training_bias(
     name_from_base,
@@ -499,21 +511,30 @@ def test_post_training_bias(
         )
 
 
-def _run_test_shap(
+def _run_test_explain(
     name_from_base,
     clarify_processor,
     clarify_processor_with_job_name_prefix,
     data_config,
     model_config,
     shap_config,
+    pdp_config,
     model_scores,
     expected_predictor_config,
 ):
     with patch.object(SageMakerClarifyProcessor, "_run", return_value=None) as mock_method:
+        explanation_configs = None
+        if shap_config and pdp_config:
+            explanation_configs = [shap_config, pdp_config]
+        elif shap_config:
+            explanation_configs = shap_config
+        elif pdp_config:
+            explanation_configs = pdp_config
+
         clarify_processor.run_explainability(
             data_config,
             model_config,
-            shap_config,
+            explanation_configs,
             model_scores=model_scores,
             wait=True,
             job_name="test",
@@ -528,23 +549,26 @@ def _run_test_shap(
                 "F3",
             ],
             "label": "Label",
-            "methods": {
-                "shap": {
-                    "baseline": [
-                        [
-                            0.26124998927116394,
-                            0.2824999988079071,
-                            0.06875000149011612,
-                        ]
-                    ],
-                    "num_samples": 100,
-                    "agg_method": "mean_sq",
-                    "use_logit": False,
-                    "save_local_shap_values": True,
-                }
-            },
             "predictor": expected_predictor_config,
         }
+        expected_explanation_configs = {}
+        if shap_config:
+            expected_explanation_configs["shap"] = {
+                "baseline": [
+                    [
+                        0.26124998927116394,
+                        0.2824999988079071,
+                        0.06875000149011612,
+                    ]
+                ],
+                "num_samples": 100,
+                "agg_method": "mean_sq",
+                "use_logit": False,
+                "save_local_shap_values": True,
+            }
+        if pdp_config:
+            expected_explanation_configs["pdp"] = {"features": ["F1", "F2"], "grid_resolution": 20}
+        expected_analysis_config["methods"] = expected_explanation_configs
         mock_method.assert_called_with(
             data_config,
             expected_analysis_config,
@@ -557,7 +581,7 @@ def _run_test_shap(
         clarify_processor_with_job_name_prefix.run_explainability(
             data_config,
             model_config,
-            shap_config,
+            explanation_configs,
             model_scores=model_scores,
             wait=True,
             experiment_config={"ExperimentName": "AnExperiment"},
@@ -574,6 +598,34 @@ def _run_test_shap(
         )
 
 
+@patch("sagemaker.utils.name_from_base", return_value=JOB_NAME)
+def test_pdp(
+    name_from_base,
+    clarify_processor,
+    clarify_processor_with_job_name_prefix,
+    data_config,
+    model_config,
+    shap_config,
+    pdp_config,
+):
+    expected_predictor_config = {
+        "model_name": "xgboost-model",
+        "instance_type": "ml.c5.xlarge",
+        "initial_instance_count": 1,
+    }
+    _run_test_explain(
+        name_from_base,
+        clarify_processor,
+        clarify_processor_with_job_name_prefix,
+        data_config,
+        model_config,
+        None,
+        pdp_config,
+        None,
+        expected_predictor_config,
+    )
+
+
 @patch("sagemaker.utils.name_from_base", return_value=JOB_NAME)
 def test_shap(
     name_from_base,
@@ -588,14 +640,15 @@ def test_shap(
         "instance_type": "ml.c5.xlarge",
         "initial_instance_count": 1,
     }
-    _run_test_shap(
+    _run_test_explain(
         name_from_base,
         clarify_processor,
         clarify_processor_with_job_name_prefix,
         data_config,
         model_config,
         shap_config,
         None,
+        None,
         expected_predictor_config,
     )
 
@@ -608,6 +661,7 @@ def test_shap_with_predicted_label(
     data_config,
     model_config,
     shap_config,
+    pdp_config,
 ):
     probability = "pr"
     label_headers = ["success"]
@@ -622,13 +676,14 @@ def test_shap_with_predicted_label(
         "probability": probability,
         "label_headers": label_headers,
     }
-    _run_test_shap(
+    _run_test_explain(
         name_from_base,
         clarify_processor,
         clarify_processor_with_job_name_prefix,
         data_config,
         model_config,
         shap_config,
+        pdp_config,
         model_scores,
         expected_predictor_config,
     )