add argument memory_tolerance to enable more flexible memory control

Author: jlyang1990

README.md

@@ -58,7 +58,7 @@ brew install libomp
 
 ## Usage
 
-The following screenshot shows a typical use case of FastTreeSHAP on [Census Income Data](https://archive.ics.uci.edu/ml/datasets/census+income). Note that the usage of FastTreeSHAP is exactly the same as the usage of [SHAP](https://github.com/slundberg/shap), except for three additional arguments in the class `TreeExplainer`: `algorithm`, `n_jobs`, and `shortcut`.
+The following screenshot shows a typical use case of FastTreeSHAP on [Census Income Data](https://archive.ics.uci.edu/ml/datasets/census+income). Note that the usage of FastTreeSHAP is exactly the same as the usage of [SHAP](https://github.com/slundberg/shap), except for four additional arguments in the class `TreeExplainer`: `algorithm`, `n_jobs`, `memory_tolerance`, and `shortcut`.
 
 `algorithm`: This argument specifies the TreeSHAP algorithm used to run FastTreeSHAP. It can take values `"v0"`, `"v1"`, `"v2"` or `"auto"`, and its default value is `"auto"`:
 * `"v0"`: Original TreeSHAP algorithm in [SHAP](https://github.com/slundberg/shap) package.
@@ -68,6 +68,8 @@ The following screenshot shows a typical use case of FastTreeSHAP on [Census Inc
 
 `n_jobs`: This argument specifies the number of parallel threads used to run FastTreeSHAP. It can take values `-1` or a positive integer. Its default value is `-1`, which means utilizing all available cores in parallel computing.
 
+`memory_tolerance`: This argument specifies the upper limit of memory allocation (in GB) to run FastTreeSHAP v2. It can take values `-1` or a positive number. Its default value is `-1`, which means allocating a maximum of 0.25 * total memory of the machine to run FastTreeSHAP v2.
+
 `shortcut`: This argument determines whether to use the TreeSHAP algorithm embedded in [XGBoost](https://github.com/dmlc/xgboost), [LightGBM](https://github.com/microsoft/LightGBM), and [CatBoost](https://github.com/catboost/catboost) packages directly when computing SHAP values for XGBoost, LightGBM, and CatBoost models and when computing SHAP interaction values for XGBoost models. Its default value is `False`, which means bypassing the "shortcut" and using the code in FastTreeSHAP package directly to compute SHAP values for XGBoost, LightGBM, and CatBoost models. Note that currently `shortcut` is automaticaly set to be True for CatBoost model, as we are working on CatBoost component in FastTreeSHAP package. More details of the usage of "shortcut" can be found in the notebooks [Census Income](notebooks/FastTreeSHAP_Census_Income.ipynb), [Superconductor](notebooks/FastTreeSHAP_Superconductor.ipynb), and [Crop Mapping](notebooks/FastTreeSHAP_Crop_Mapping.ipynb).
 
 ![FastTreeSHAP Adult Screenshot1](docs/images/fasttreeshap_adult_screenshot1.png)

fasttreeshap/__init__.py

@@ -3,7 +3,7 @@
 import warnings
 import sys
 
-__version__ = '0.1.1'
+__version__ = '0.1.2'
 
 # check python version
 if (sys.version_info < (3, 0)):

fasttreeshap/explainers/_tree.py

@@ -60,7 +60,7 @@ class Tree(Explainer):
     implementations either inside an externel model package or in the local compiled C extention.
     """
 
-    def __init__(self, model, data = None, model_output="raw", feature_perturbation="interventional", algorithm="auto", n_jobs=-1, feature_names=None, approximate=False, shortcut=False, **deprecated_options):
+    def __init__(self, model, data = None, model_output="raw", feature_perturbation="tree_path_dependent", algorithm="auto", n_jobs=-1, memory_tolerance=-1, feature_names=None, approximate=False, shortcut=False, **deprecated_options):
         """ Build a new Tree explainer for the passed model.
 
         Parameters
@@ -99,6 +99,10 @@ def __init__(self, model, data = None, model_output="raw", feature_perturbation=
             all available cores in parallel computing (Setting OMP_NUM_THREADS is unnecessary since n_jobs will
             overwrite this parameter).
 
+        memory_tolerance : -1 (default), or a positive number
+            Upper limit of memory allocation (in GB) to run Fast TreeSHAP v2. The default value of memory_tolerance is -1,
+            which allocates a maximum of 0.25 * total memory of the machine to run Fast TreeSHAP v2.
+
         model_output : "raw", "probability", "log_loss", or model method name
             What output of the model should be explained. If "raw" then we explain the raw output of the
             trees, which varies by model. For regression models "raw" is the standard output, for binary
@@ -182,6 +186,7 @@ def __init__(self, model, data = None, model_output="raw", feature_perturbation=
             self.n_jobs = 1
         else:
             self.n_jobs = min(int(n_jobs), os.cpu_count())
+        self.memory_tolerance = memory_tolerance
         self.expected_value = None
         self.model = TreeEnsemble(model, self.data, self.data_missing, model_output)
         self.model_output = model_output
@@ -490,6 +495,8 @@ def _memory_check(self, X):
             memory_tolerance = 0.25 * psutil.virtual_memory().total
         except:
             memory_tolerance = 4294967296  # 4GB
+        if self.memory_tolerance > 0:
+            memory_tolerance = min(memory_tolerance, self.memory_tolerance * 1073741824)
         return memory_usage_1 <= memory_tolerance, memory_usage_2 <= memory_tolerance