Add Occlusion to Insights

Author: edward-io

captum/insights/api.py

@@ -89,7 +89,7 @@ class FilterConfig(NamedTuple):
         arg: config.value  # type: ignore
         for arg, config in ATTRIBUTION_METHOD_CONFIG[
             IntegratedGradients.get_name()
-        ].items()
+        ].params.items()
     }
     prediction: str = "all"
     classes: List[str] = []
@@ -221,6 +221,12 @@ def _calculate_attribution(
         attribution_cls = ATTRIBUTION_NAMES_TO_METHODS[self._config.attribution_method]
         attribution_method = attribution_cls(net)
         args = self._config.attribution_arguments
+        param_config = ATTRIBUTION_METHOD_CONFIG[self._config.attribution_method]
+        if param_config.post_process:
+            for k, v in args.items():
+                if k in param_config.post_process:
+                    args[k] = param_config.post_process[k](v)
+
         # TODO support multiple baselines
         baseline = baselines[0] if baselines and len(baselines) > 0 else None
         label = (
@@ -542,6 +548,8 @@ def get_insights_config(self):
         return {
             "classes": self.classes,
             "methods": list(ATTRIBUTION_NAMES_TO_METHODS.keys()),
-            "method_arguments": namedtuple_to_dict(ATTRIBUTION_METHOD_CONFIG),
+            "method_arguments": namedtuple_to_dict(
+                {k: v.params for (k, v) in ATTRIBUTION_METHOD_CONFIG.items()}
+            ),
             "selected_method": self._config.attribution_method,
         }

captum/insights/config.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-from typing import Dict, List, NamedTuple, Optional, Tuple
+from typing import Dict, List, NamedTuple, Optional, Tuple, Callable, Any, Union
 
 from captum.attr import (
     Deconvolution,
@@ -9,10 +9,14 @@
     InputXGradient,
     IntegratedGradients,
     Saliency,
+    Occlusion,
 )
 from captum.attr._utils.approximation_methods import SUPPORTED_METHODS
 
 
+Config = Union[NumberConfig, StrEnumConfig, StrConfig]
+
+
 class NumberConfig(NamedTuple):
     value: int = 1
     limit: Tuple[Optional[int], Optional[int]] = (None, None)
@@ -25,6 +29,11 @@ class StrEnumConfig(NamedTuple):
     type: str = "enum"
 
 
+class StrConfig(NamedTuple):
+    value: str
+    type: str = "string"
+
+
 SUPPORTED_ATTRIBUTION_METHODS = [
     Deconvolution,
     DeepLift,
@@ -33,20 +42,47 @@ class StrEnumConfig(NamedTuple):
     IntegratedGradients,
     Saliency,
     FeatureAblation,
+    Occlusion,
 ]
 
+
+class ConfigParameters(NamedTuple):
+    params: Dict[str, Config]
+    help_info: Optional[str] = None  # TODO fill out help for each method
+    post_process: Optional[Dict[str, Callable[[Any], Any]]] = None
+
+
 ATTRIBUTION_NAMES_TO_METHODS = {
     # mypy bug - treating it as a type instead of a class
     cls.get_name(): cls  # type: ignore
     for cls in SUPPORTED_ATTRIBUTION_METHODS
 }
 
-ATTRIBUTION_METHOD_CONFIG: Dict[str, Dict[str, tuple]] = {
-    IntegratedGradients.get_name(): {
-        "n_steps": NumberConfig(value=25, limit=(2, None)),
-        "method": StrEnumConfig(limit=SUPPORTED_METHODS, value="gausslegendre"),
-    },
-    FeatureAblation.get_name(): {
-        "perturbations_per_eval": NumberConfig(value=1, limit=(1, 100)),
-    },
+
+def _str_to_tuple(s):
+    if isinstance(s, tuple):
+        return s
+    return tuple([int(i) for i in s.split()])
+
+
+ATTRIBUTION_METHOD_CONFIG: Dict[str, ConfigParameters] = {
+    IntegratedGradients.get_name(): ConfigParameters(
+        params={
+            "n_steps": NumberConfig(value=25, limit=(2, None)),
+            "method": StrEnumConfig(limit=SUPPORTED_METHODS, value="gausslegendre"),
+        }
+    ),
+    FeatureAblation.get_name(): ConfigParameters(
+        params={"perturbations_per_eval": NumberConfig(value=1, limit=(1, 100))}
+    ),
+    Occlusion.get_name(): ConfigParameters(
+        params={
+            "sliding_window_shapes": StrConfig(value=""),
+            "strides": StrConfig(value=""),
+        },
+        post_process={
+            "sliding_window_shapes": _str_to_tuple,
+            "strides": _str_to_tuple,
+        },
+    ),
 }

captum/insights/frontend/src/App.js

@@ -8,6 +8,7 @@ import "./App.css";
 const ConfigType = Object.freeze({
   Number: "number",
   Enum: "enum",
+  String: "string",
 });
 
 const Plot = createPlotlyComponent(Plotly);
@@ -153,62 +154,71 @@ class FilterContainer extends React.Component {
   }
 }
 
-class ClassFilter extends React.Component {
-  render() {
-    return (
-      <ReactTags
-        tags={this.props.classes}
-        autofocus={false}
-        suggestions={this.props.suggestedClasses}
-        handleDelete={this.props.handleClassDelete}
-        handleAddition={this.props.handleClassAdd}
-        minQueryLength={0}
-        placeholder="add new class..."
+function ClassFilter(props) {
+  return (
+    <ReactTags
+      tags={props.classes}
+      autofocus={false}
+      suggestions={props.suggestedClasses}
+      handleDelete={props.handleClassDelete}
+      handleAddition={props.handleClassAdd}
+      minQueryLength={0}
+      placeholder="add new class..."
+    />
+  );
+}
+
+function NumberArgument(props) {
+  var min = props.limit[0];
+  var max = props.limit[1];
+  return (
+    <div>
+      {props.name}:
+      <input
+        className={cx([styles.input, styles["input--narrow"]])}
+        name={props.name}
+        type="number"
+        value={props.value}
+        min={min}
+        max={max}
+        onChange={props.handleInputChange}
       />
-    );
-  }
+    </div>
+  );
 }
 
-class NumberArgument extends React.Component {
-  render() {
-    var min = this.props.limit[0];
-    var max = this.props.limit[1];
-    return (
-      <div>
-        {this.props.name + ": "}
-        <input
-          className={cx([styles.input, styles["input--narrow"]])}
-          name={this.props.name}
-          type="number"
-          value={this.props.value}
-          min={min}
-          max={max}
-          onChange={this.props.handleInputChange}
-        />
-      </div>
-    );
-  }
+function EnumArgument(props) {
+  const options = props.limit.map((item, key) => (
+    <option value={item}>{item}</option>
+  ));
+  return (
+    <div>
+      {props.name}:
+      <select
+        className={styles.select}
+        name={props.name}
+        value={props.value}
+        onChange={props.handleInputChange}
+      >
+        {options}
+      </select>
+    </div>
+  );
 }
 
-class EnumArgument extends React.Component {
-  render() {
-    const options = this.props.limit.map((item, key) => (
-      <option value={item}>{item}</option>
-    ));
-    return (
-      <div>
-        {this.props.name + ": "}
-        <select
-          className={styles.select}
-          name={this.props.name}
-          value={this.props.value}
-          onChange={this.props.handleInputChange}
-        >
-          {options}
-        </select>
-      </div>
-    );
-  }
+function StringArgument(props) {
+  return (
+    <div>
+      {props.name}:
+      <input
+        className={cx([styles.input, styles["input--narrow"]])}
+        name={props.name}
+        type="text"
+        value={props.value}
+        onChange={props.handleInputChange}
+      />
+    </div>
+  );
 }
 
 class Filter extends React.Component {
@@ -232,6 +242,14 @@ class Filter extends React.Component {
             handleInputChange={this.props.handleArgumentChange}
           />
         );
+      case ConfigType.String:
+        return (
+          <StringArgument
+            name={name}
+            value={config.value}
+            handleInputChange={this.props.handleArgumentChange}
+          />
+        );
     }
   };