wehub-resource-sync
265 lines
8.9 KiB
Python
265 lines
8.9 KiB
Python
import os
|
|
from typing import Any, NamedTuple
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
import shap
|
|
from sklearn.datasets import load_diabetes, load_iris
|
|
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
|
|
|
|
import mlflow
|
|
from mlflow import MlflowClient
|
|
from mlflow.utils.file_utils import local_file_uri_to_path
|
|
|
|
|
|
class ModelWithExplanation(NamedTuple):
|
|
model: Any
|
|
X: Any
|
|
shap_values: Any
|
|
base_values: Any
|
|
|
|
|
|
def yield_artifacts(run_id, path=None):
|
|
"""
|
|
Yields all artifacts in the specified run.
|
|
"""
|
|
client = MlflowClient()
|
|
for item in client.list_artifacts(run_id, path):
|
|
if item.is_dir:
|
|
yield from yield_artifacts(run_id, item.path)
|
|
else:
|
|
yield item.path
|
|
|
|
|
|
def get_iris():
|
|
data = load_iris()
|
|
return (
|
|
pd.DataFrame(data.data, columns=data.feature_names),
|
|
pd.Series(data.target, name="target"),
|
|
)
|
|
|
|
|
|
def get_diabetes():
|
|
X, y = load_diabetes(return_X_y=True, as_frame=True)
|
|
return X.iloc[:100, :4], y.iloc[:100]
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def regressor():
|
|
X, y = get_diabetes()
|
|
model = RandomForestRegressor()
|
|
model.fit(X, y)
|
|
|
|
explainer = shap.KernelExplainer(model.predict, shap.kmeans(X, 100))
|
|
shap_values = explainer.shap_values(X)
|
|
|
|
return ModelWithExplanation(model, X, shap_values, explainer.expected_value)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def classifier():
|
|
X, y = get_iris()
|
|
model = RandomForestClassifier()
|
|
model.fit(X, y)
|
|
|
|
explainer = shap.KernelExplainer(model.predict_proba, shap.kmeans(X, 100))
|
|
shap_values = explainer.shap_values(X)
|
|
|
|
return ModelWithExplanation(model, X, shap_values, explainer.expected_value)
|
|
|
|
|
|
@pytest.mark.parametrize("np_obj", [float(0.0), np.array([0.0])])
|
|
def test_log_numpy(np_obj):
|
|
with mlflow.start_run() as run:
|
|
mlflow.shap._log_numpy(np_obj, "test.npy")
|
|
mlflow.shap._log_numpy(np_obj, "test.npy", artifact_path="dir")
|
|
|
|
artifacts = set(yield_artifacts(run.info.run_id))
|
|
assert artifacts == {"test.npy", "dir/test.npy"}
|
|
|
|
|
|
def test_log_matplotlib_figure():
|
|
fig, ax = plt.subplots()
|
|
ax.plot([0, 1], [2, 3])
|
|
|
|
with mlflow.start_run() as run:
|
|
mlflow.shap._log_matplotlib_figure(fig, "test.png")
|
|
mlflow.shap._log_matplotlib_figure(fig, "test.png", artifact_path="dir")
|
|
|
|
artifacts = set(yield_artifacts(run.info.run_id))
|
|
assert artifacts == {"test.png", "dir/test.png"}
|
|
|
|
|
|
def test_log_explanation_with_regressor(regressor):
|
|
model = regressor.model
|
|
X = regressor.X
|
|
|
|
with mlflow.start_run() as run:
|
|
explanation_path = mlflow.shap.log_explanation(model.predict, X)
|
|
|
|
# Assert no figure is open
|
|
assert len(plt.get_fignums()) == 0
|
|
|
|
artifact_path = "model_explanations_shap"
|
|
artifacts = set(yield_artifacts(run.info.run_id))
|
|
|
|
assert explanation_path == os.path.join(run.info.artifact_uri, artifact_path)
|
|
assert artifacts == {
|
|
os.path.join(artifact_path, "base_values.npy"),
|
|
os.path.join(artifact_path, "shap_values.npy"),
|
|
os.path.join(artifact_path, "summary_bar_plot.png"),
|
|
}
|
|
|
|
local_path = local_file_uri_to_path(explanation_path)
|
|
shap_values = np.load(os.path.join(local_path, "shap_values.npy"))
|
|
base_values = np.load(os.path.join(local_path, "base_values.npy"))
|
|
np.testing.assert_array_equal(shap_values, regressor.shap_values)
|
|
np.testing.assert_array_equal(base_values, regressor.base_values)
|
|
|
|
|
|
def test_log_explanation_with_classifier(classifier):
|
|
model = classifier.model
|
|
X = classifier.X
|
|
|
|
with mlflow.start_run() as run:
|
|
explanation_uri = mlflow.shap.log_explanation(model.predict_proba, X)
|
|
|
|
# Assert no figure is open
|
|
assert len(plt.get_fignums()) == 0
|
|
|
|
artifact_path = "model_explanations_shap"
|
|
artifacts = set(yield_artifacts(run.info.run_id))
|
|
|
|
assert explanation_uri == os.path.join(run.info.artifact_uri, artifact_path)
|
|
assert artifacts == {
|
|
os.path.join(artifact_path, "base_values.npy"),
|
|
os.path.join(artifact_path, "shap_values.npy"),
|
|
os.path.join(artifact_path, "summary_bar_plot.png"),
|
|
}
|
|
|
|
local_path = local_file_uri_to_path(explanation_uri)
|
|
shap_values = np.load(os.path.join(local_path, "shap_values.npy"))
|
|
base_values = np.load(os.path.join(local_path, "base_values.npy"))
|
|
np.testing.assert_array_equal(shap_values, classifier.shap_values)
|
|
np.testing.assert_array_equal(base_values, classifier.base_values)
|
|
|
|
|
|
@pytest.mark.parametrize("artifact_path", ["dir", "dir1/dir2"])
|
|
def test_log_explanation_with_artifact_path(regressor, artifact_path):
|
|
model = regressor.model
|
|
X = regressor.X
|
|
|
|
with mlflow.start_run() as run:
|
|
explanation_path = mlflow.shap.log_explanation(model.predict, X, artifact_path)
|
|
|
|
# Assert no figure is open
|
|
assert len(plt.get_fignums()) == 0
|
|
|
|
artifacts = set(yield_artifacts(run.info.run_id))
|
|
|
|
assert explanation_path == os.path.join(run.info.artifact_uri, artifact_path)
|
|
assert artifacts == {
|
|
os.path.join(artifact_path, "base_values.npy"),
|
|
os.path.join(artifact_path, "shap_values.npy"),
|
|
os.path.join(artifact_path, "summary_bar_plot.png"),
|
|
}
|
|
|
|
local_path = local_file_uri_to_path(explanation_path)
|
|
shap_values = np.load(os.path.join(local_path, "shap_values.npy"))
|
|
base_values = np.load(os.path.join(local_path, "base_values.npy"))
|
|
np.testing.assert_array_equal(shap_values, regressor.shap_values)
|
|
np.testing.assert_array_equal(base_values, regressor.base_values)
|
|
|
|
|
|
def test_log_explanation_without_active_run(regressor):
|
|
model = regressor.model
|
|
X = regressor.X.values
|
|
|
|
with mlflow.start_run() as run:
|
|
explanation_uri = mlflow.shap.log_explanation(model.predict, X)
|
|
|
|
# Assert no figure is open
|
|
assert len(plt.get_fignums()) == 0
|
|
|
|
artifact_path = "model_explanations_shap"
|
|
artifacts = set(yield_artifacts(run.info.run_id))
|
|
|
|
assert explanation_uri == os.path.join(run.info.artifact_uri, artifact_path)
|
|
assert artifacts == {
|
|
os.path.join(artifact_path, "base_values.npy"),
|
|
os.path.join(artifact_path, "shap_values.npy"),
|
|
os.path.join(artifact_path, "summary_bar_plot.png"),
|
|
}
|
|
|
|
local_path = local_file_uri_to_path(explanation_uri)
|
|
shap_values = np.load(os.path.join(local_path, "shap_values.npy"))
|
|
base_values = np.load(os.path.join(local_path, "base_values.npy"))
|
|
np.testing.assert_array_equal(shap_values, regressor.shap_values)
|
|
np.testing.assert_array_equal(base_values, regressor.base_values)
|
|
|
|
|
|
def test_log_explanation_with_numpy_array(regressor):
|
|
model = regressor.model
|
|
X = regressor.X.values
|
|
|
|
with mlflow.start_run() as run:
|
|
explanation_uri = mlflow.shap.log_explanation(model.predict, X)
|
|
|
|
# Assert no figure is open
|
|
assert len(plt.get_fignums()) == 0
|
|
|
|
artifact_path = "model_explanations_shap"
|
|
artifacts = set(yield_artifacts(run.info.run_id))
|
|
|
|
assert explanation_uri == os.path.join(run.info.artifact_uri, artifact_path)
|
|
assert artifacts == {
|
|
os.path.join(artifact_path, "base_values.npy"),
|
|
os.path.join(artifact_path, "shap_values.npy"),
|
|
os.path.join(artifact_path, "summary_bar_plot.png"),
|
|
}
|
|
|
|
local_path = local_file_uri_to_path(explanation_uri)
|
|
shap_values = np.load(os.path.join(local_path, "shap_values.npy"))
|
|
base_values = np.load(os.path.join(local_path, "base_values.npy"))
|
|
np.testing.assert_array_equal(shap_values, regressor.shap_values)
|
|
np.testing.assert_array_equal(base_values, regressor.base_values)
|
|
|
|
|
|
def test_log_explanation_with_small_features():
|
|
"""
|
|
Verifies that `log_explanation` does not fail even when `features` has less records than
|
|
`_MAXIMUM_BACKGROUND_DATA_SIZE`.
|
|
"""
|
|
num_rows = 50
|
|
assert num_rows < mlflow.shap._MAXIMUM_BACKGROUND_DATA_SIZE
|
|
|
|
X, y = get_diabetes()
|
|
X = X.iloc[:num_rows]
|
|
y = y[:num_rows]
|
|
model = RandomForestRegressor()
|
|
model.fit(X, y)
|
|
|
|
with mlflow.start_run() as run:
|
|
explanation_uri = mlflow.shap.log_explanation(model.predict, X)
|
|
|
|
artifact_path = "model_explanations_shap"
|
|
artifacts = set(yield_artifacts(run.info.run_id))
|
|
|
|
assert explanation_uri == os.path.join(run.info.artifact_uri, artifact_path)
|
|
assert artifacts == {
|
|
os.path.join(artifact_path, "base_values.npy"),
|
|
os.path.join(artifact_path, "shap_values.npy"),
|
|
os.path.join(artifact_path, "summary_bar_plot.png"),
|
|
}
|
|
|
|
explainer = shap.KernelExplainer(model.predict, shap.kmeans(X, num_rows))
|
|
shap_values_expected = explainer.shap_values(X)
|
|
|
|
local_path = local_file_uri_to_path(explanation_uri)
|
|
base_values = np.load(os.path.join(local_path, "base_values.npy"))
|
|
shap_values = np.load(os.path.join(local_path, "shap_values.npy"))
|
|
np.testing.assert_array_equal(base_values, explainer.expected_value)
|
|
np.testing.assert_array_equal(shap_values, shap_values_expected)
|