Small bug fixes

Karol-G · Jun 20, 2020 · 2f143f8 · 2f143f8
1 parent aa9e22e
commit 2f143f8
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/gcam/backends/grad_cam.py b/gcam/backends/grad_cam.py
@@ -91,15 +91,15 @@ def generate(self):
         if self._target_layers == "auto":
             layer, fmaps, grads = self._auto_layer_selection()
             self._check_hooks(layer)
-            attention_map = self._generate_helper(fmaps, grads).cpu().numpy()
+            attention_map = self._generate_helper(fmaps, grads, layer).cpu().numpy()
             attention_maps = {layer: attention_map}
         else:
             for layer in self.target_layers:
                 self._check_hooks(layer)
                 if self.registered_hooks[layer][0] and self.registered_hooks[layer][1]:
                     fmaps = self._find(self.fmap_pool, layer)
                     grads = self._find(self.grad_pool, layer)
-                    attention_map = self._generate_helper(fmaps, grads)
+                    attention_map = self._generate_helper(fmaps, grads, layer)
                     attention_maps[layer] = attention_map.cpu().numpy()
         if not attention_maps:
             raise ValueError("None of the hooks registered to the target layers")
@@ -153,7 +153,7 @@ def _compute_grad_weights(self, grads):
         else:
             return F.adaptive_avg_pool3d(grads, 1)
 
-    def _generate_helper(self, fmaps, grads):
+    def _generate_helper(self, fmaps, grads, layer):
         weights = self._compute_grad_weights(grads)
         attention_map = torch.mul(fmaps, weights)
         B, _, *data_shape = attention_map.shape

diff --git a/gcam/backends/grad_cam_pp.py b/gcam/backends/grad_cam_pp.py
@@ -14,7 +14,7 @@ def __init__(self, model, target_layers=None, postprocessor=None, retain_graph=F
         """
         super(GradCamPP, self).__init__(model, target_layers=target_layers, postprocessor=postprocessor, retain_graph=retain_graph)
 
-    def _generate_helper(self, fmaps, grads):
+    def _generate_helper(self, fmaps, grads, layer):
         B, C, *data_shape = grads.size()
 
         alpha_num = grads.pow(2)

diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="gcam",
-    version="0.0.13",
+    version="0.0.14",
     author="Karol Gotkowski",
     author_email="KarolGotkowski@gmx.de",
     description="An easy to use framework that makes model predictions more interpretable for humans.",