Style changes and minor fixes to Insights (#68)

* Style changes and minor fixes to Insights

* fix tests and suppress warning

Author: edward-io

captum/insights/api.py

@@ -62,7 +62,7 @@ def _calculate_attribution(
     ) -> Tensor:
         ig = IntegratedGradients(net)
         # TODO support multiple baselines
-        label = None if label is None or len(label.shape) == 0 else label
+        label = None if label is None or label.nelement() == 0 else label
         attr_ig, _ = ig.attribute(
             data,
             baselines=baselines[0],
@@ -124,7 +124,7 @@ def _calculate_net_contrib(self, attrs_per_input_feature: List[Tensor]):
         if norm > 0:
             net_contrib /= norm
 
-        return net_contrib
+        return net_contrib.tolist()
 
     def visualize(self) -> List[VisualizationOutput]:
         batch_data = next(self.dataset)
@@ -169,11 +169,11 @@ def visualize(self) -> List[VisualizationOutput]:
 
             label = batch_data.labels[i]
 
-            if len(outputs) == 1:
+            if outputs.nelement() == 1:
                 scores = outputs
                 predicted = scores.round().to(torch.int)
             else:
-                scores, predicted = outputs.topk(min(4, len(outputs)))
+                scores, predicted = outputs.topk(min(4, outputs.shape[-1]))
 
             scores = scores.cpu().squeeze(0)
             predicted = predicted.cpu().squeeze(0)

captum/insights/features.py

@@ -1,4 +1,5 @@
 import base64
+import warnings
 from collections import namedtuple
 from io import BytesIO
 from typing import Callable, List, Optional, Union
@@ -12,7 +13,10 @@
 
 def _convert_figure_base64(fig):
     buff = BytesIO()
-    fig.savefig(buff, format="png", pad_inches=0.0)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        fig.tight_layout()  # removes padding
+    fig.savefig(buff, format="png")
     base64img = base64.b64encode(buff.getvalue()).decode("utf-8")
     return base64img
 

captum/insights/frontend/src/App.css

@@ -170,16 +170,27 @@ button {
 }
 
 .gallery__item__image img {
-  height: 130px;
+  height: 200px;
   width: auto;
 }
 
+.gallery__item__description {
+  text-align: center;
+}
+
 .bar-chart__group {
-  padding: 8px 0;
+  padding: 2px 0;
+  display: flex;
 }
+
 .bar-chart__group__bar {
-  height: 10px;
-  border-radius: 4px;
+  width: 10px;
+  border-radius: 2px;
+  flex-shrink: 0;
+}
+
+.bar-chart__group__title {
+  padding-left: 8px;
 }
 
 .percentage-blue {

captum/insights/frontend/src/App.js

@@ -160,7 +160,9 @@ function ImageFeature(props) {
             <div className="gallery__item__image">
               <img src={"data:image/png;base64," + props.data.modified} />
             </div>
-            <div className="gallery__item__description">Gradient Overlay</div>
+            <div className="gallery__item__description">
+              Attribution Magnitude
+            </div>
           </div>
         </div>
       </div>
@@ -206,18 +208,24 @@ function Feature(props) {
 
 class Contributions extends React.Component {
   render() {
-    return this.props.feature_outputs.map(f => (
-      <div className="bar-chart__group">
-        <div
-          className={cx({
-            "bar-chart__group__bar": true,
-            [getPercentageColor(f.contribution)]: true
-          })}
-          width={f.contribution + "%"}
-        />
-        <div className="bar-chart__group__title">{f.name}</div>
-      </div>
-    ));
+    return this.props.feature_outputs.map(f => {
+      // pad bar height so features with 0 contribution can still be seen
+      // in graph
+      const contribution = f.contribution * 100;
+      const bar_height = contribution > 10 ? contribution : contribution + 10;
+      return (
+        <div className="bar-chart__group">
+          <div
+            className={cx([
+              "bar-chart__group__bar",
+              getPercentageColor(contribution)
+            ])}
+            style={{ height: bar_height + "px" }}
+          />
+          <div className="bar-chart__group__title">{f.name}</div>
+        </div>
+      );
+    });
   }
 }
 

tests/insights/test_contribution.py

@@ -1,14 +1,15 @@
 from __future__ import print_function
-from typing import Callable, List, Optional, Union
 
-import torch
-import torch.nn as nn
+import unittest
+from typing import Callable, List, Optional, Union
 
 from captum.insights.api import AttributionVisualizer, Data
-from captum.insights.features import ImageFeature, BaseFeature, FeatureOutput
-
+from captum.insights.features import BaseFeature, FeatureOutput, ImageFeature
 from tests.attr.helpers.utils import BaseTest
 
+import torch
+import torch.nn as nn
+
 
 class RealFeature(BaseFeature):
     def __init__(
@@ -97,15 +98,15 @@ def forward(self, img, misc):
 
 
 def _labelled_img_data(num_samples=10, width=8, height=8, depth=3, num_labels=10):
-    for i in range(num_samples):
+    for _ in range(num_samples):
         yield torch.empty(depth, height, width).uniform_(0, 1), torch.randint(
             num_labels, (1,)
         )
 
 
 def _multi_modal_data(img_dataset, feature_size=256):
     def misc_data(length, feature_size=None):
-        for i in range(length):
+        for _ in range(length):
             yield torch.randn(feature_size)
 
     misc_dataset = misc_data(length=len(img_dataset), feature_size=feature_size)
@@ -165,11 +166,8 @@ def test_one_feature(self):
         outputs = visualizer.visualize()
 
         for output in outputs:
-            contribs = torch.stack(
-                [feature.contribution for feature in output.feature_outputs]
-            )
-            total_contrib = torch.sum(torch.abs(contribs))
-            self.assertAlmostEqual(total_contrib.item(), 1.0, places=6)
+            total_contrib = sum(abs(f.contribution) for f in output.feature_outputs)
+            self.assertAlmostEqual(total_contrib, 1.0, places=6)
 
     def test_multi_features(self):
         batch_size = 2
@@ -183,7 +181,7 @@ def test_multi_features(self):
             img_dataset=img_dataset, feature_size=misc_feature_size
         )
         # NOTE: using DataLoader to batch the inputs since
-        # AttributionVisualizer requires the input to be of size `B x ...`
+        # AttributionVisualizer requires the input to be of size `N x ...`
         data_loader = torch.utils.data.DataLoader(
             list(dataset), batch_size=batch_size, shuffle=False, num_workers=0
         )
@@ -211,13 +209,14 @@ def test_multi_features(self):
         outputs = visualizer.visualize()
 
         for output in outputs:
-            contribs = torch.stack(
-                [feature.contribution for feature in output.feature_outputs]
-            )
-            total_contrib = torch.sum(torch.abs(contribs))
-            self.assertAlmostEqual(total_contrib.item(), 1.0, places=6)
+            total_contrib = sum(abs(f.contribution) for f in output.feature_outputs)
+            self.assertAlmostEqual(total_contrib, 1.0, places=6)
 
     # TODO: add test for multiple models (related to TODO in captum/insights/api.py)
     #
     # TODO: add test to make the attribs == 0 -- error occurs
     #       I know (through manual testing) that this breaks some existing code
+
+
+if __name__ == "__main__":
+    unittest.main()