update a simpler version

Author: cai2-huaiguang

pytorch_grad_cam/activations_and_gradients.py

@@ -2,11 +2,12 @@ class ActivationsAndGradients:
     """ Class for extracting activations and
     registering gradients from targetted intermediate layers """
 
-    def __init__(self, model, target_layers, reshape_transform):
+    def __init__(self, model, target_layers, reshape_transform, detach=True):
         self.model = model
         self.gradients = []
         self.activations = []
         self.reshape_transform = reshape_transform
+        self.detach = detach
         self.handles = []
         for target_layer in target_layers:
             self.handles.append(
@@ -18,10 +19,12 @@ def __init__(self, model, target_layers, reshape_transform):
 
     def save_activation(self, module, input, output):
         activation = output
-
-        if self.reshape_transform is not None:
-            activation = self.reshape_transform(activation)
-        self.activations.append(activation.cpu().detach())
+        if self.detach:
+            if self.reshape_transform is not None:
+                activation = self.reshape_transform(activation)
+            self.activations.append(activation.cpu().detach())
+        else:
+            self.activations.append(activation)
 
     def save_gradient(self, module, input, output):
         if not hasattr(output, "requires_grad") or not output.requires_grad:
@@ -30,9 +33,12 @@ def save_gradient(self, module, input, output):
 
         # Gradients are computed in reverse order
         def _store_grad(grad):
-            if self.reshape_transform is not None:
-                grad = self.reshape_transform(grad)
-            self.gradients = [grad.cpu().detach()] + self.gradients
+            if self.detach:
+                if self.reshape_transform is not None:
+                    grad = self.reshape_transform(grad)
+                self.gradients = [grad.cpu().detach()] + self.gradients
+            else:
+                self.gradients = [grad] + self.gradients
 
         output.register_hook(_store_grad)
 

pytorch_grad_cam/activations_and_gradients_no_detach.py

@@ -19,6 +19,7 @@ def __init__(
         compute_input_gradient: bool = False,
         uses_gradients: bool = True,
         tta_transforms: Optional[tta.Compose] = None,
+        detach: bool = True,
     ) -> None:
         self.model = model.eval()
         self.target_layers = target_layers
@@ -45,7 +46,8 @@ def __init__(
         else:
             self.tta_transforms = tta_transforms
 
-        self.activations_and_grads = ActivationsAndGradients(self.model, target_layers, reshape_transform)
+        self.detach = detach
+        self.activations_and_grads = ActivationsAndGradients(self.model, target_layers, reshape_transform, self.detach)
 
     """ Get a vector of weights for every channel in the target layer.
         Methods that return weights channels,
@@ -71,6 +73,8 @@ def get_cam_image(
         eigen_smooth: bool = False,
     ) -> np.ndarray:
         weights = self.get_cam_weights(input_tensor, target_layer, targets, activations, grads)
+        if isinstance(activations, torch.Tensor):
+            activations = activations.cpu().detach().numpy()
         # 2D conv
         if len(activations.shape) == 4:
             weighted_activations = weights[:, :, None, None] * activations
@@ -132,8 +136,12 @@ def get_target_width_height(self, input_tensor: torch.Tensor) -> Tuple[int, int]
     def compute_cam_per_layer(
         self, input_tensor: torch.Tensor, targets: List[torch.nn.Module], eigen_smooth: bool
     ) -> np.ndarray:
-        activations_list = [a.cpu().data.numpy() for a in self.activations_and_grads.activations]
-        grads_list = [g.cpu().data.numpy() for g in self.activations_and_grads.gradients]
+        if self.detach:
+            activations_list = [a.cpu().data.numpy() for a in self.activations_and_grads.activations]
+            grads_list = [g.cpu().data.numpy() for g in self.activations_and_grads.gradients]
+        else:
+            activations_list = [a for a in self.activations_and_grads.activations]
+            grads_list = [g for g in self.activations_and_grads.gradients]
         target_size = self.get_target_width_height(input_tensor)
 
         cam_per_target_layer = []

pytorch_grad_cam/base_cam.py

@@ -1,32 +1,22 @@
 from typing import Callable, List, Optional, Tuple
-
-import numpy as np
-import torch
 from pytorch_grad_cam.base_cam import BaseCAM
-from scipy.signal import convolve2d
-from scipy.ndimage import gaussian_filter
-import cv2
-
-from pytorch_grad_cam.activations_and_gradients_no_detach import ActivationsAndGradients_no_detach
-from pytorch_grad_cam.utils.image import scale_cam_image
-from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
-from pytorch_grad_cam.utils.svd_on_activations import get_2d_projection
+import torch
+import numpy as np
 
 """
 Weighting the activation maps using Gradient and Hessian-Vector Product.
-This method (https://arxiv.org/abs/2501.06261) reinterpret CAM methods from a Shapley value perspective.
+This method (https://arxiv.org/abs/2501.06261) reinterpret CAM methods (include GradCAM, HiResCAM and the original CAM) from a Shapley value perspective.
 """
 class ShapleyCAM(BaseCAM):
     def __init__(self, model, target_layers,
-                 reshape_transform=None):
+                 reshape_transform=None, detach=False):
         super(
             ShapleyCAM,
             self).__init__(
             model,
             target_layers,
-            reshape_transform)
-
-        self.activations_and_grads = ActivationsAndGradients_no_detach(self.model, target_layers, reshape_transform)
+            reshape_transform,
+            detach = detach)
 
     def forward(
         self, input_tensor: torch.Tensor, targets: List[torch.nn.Module], eigen_smooth: bool = False
@@ -44,6 +34,7 @@ def forward(
         if self.uses_gradients:
             self.model.zero_grad()
             loss = sum([target(output) for target, output in zip(targets, outputs)])
+            # keep the graph
             torch.autograd.grad(loss, input_tensor,  retain_graph = True, create_graph = True)
 
         # In most of the saliency attribution papers, the saliency is
@@ -65,96 +56,36 @@ def get_cam_weights(self,
                         target_category,
                         activations,
                         grads):
-        activations: List[Tensor]  # type: ignore[assignment]
-        grads: List[Tensor]  # type: ignore[assignment]
-
+        
         hvp = torch.autograd.grad(
             outputs=grads,
             inputs=activations,
             grad_outputs=activations,
             retain_graph=False,
             allow_unused=True
         )[0]
-        # print(torch.max(hvp[0]).item())  # verify that hvp is not all zeros
+        # print(torch.max(hvp[0]).item())  # Use .item() to get the scalar value
         if hvp is None:
             hvp = torch.tensor(0).to(self.device)
-        elif self.activations_and_grads.reshape_transform is not None:
-            hvp = self.activations_and_grads.reshape_transform(hvp)
+        else:
+            if self.activations_and_grads.reshape_transform is not None:
+                hvp = self.activations_and_grads.reshape_transform(hvp)
 
         if self.activations_and_grads.reshape_transform is not None:
             activations = self.activations_and_grads.reshape_transform(activations)
             grads = self.activations_and_grads.reshape_transform(grads)
-        weight = (grads  - 0.5 * hvp).cpu().detach().numpy()
-        activations = activations.cpu().detach().numpy()
-        grads = grads.cpu().detach().numpy()
-
 
+        weight = (grads  - 0.5 * hvp).detach().cpu().numpy()
         # 2D image
         if len(activations.shape) == 4:
             weight = np.mean(weight, axis=(2, 3))
-            return weight, activations
+            return weight
 
         # 3D image
         elif len(activations.shape) == 5:
             weight = np.mean(weight, axis=(2, 3, 4))
-            return weight, activations
+            return weight
 
         else:
             raise ValueError("Invalid grads shape."
                              "Shape of grads should be 4 (2D image) or 5 (3D image).")
-
-
-
-    def get_cam_image(
-        self,
-        input_tensor: torch.Tensor,
-        target_layer: torch.nn.Module,
-        targets: List[torch.nn.Module],
-        activations: torch.Tensor,
-        grads: torch.Tensor,
-        eigen_smooth: bool = False,
-    ) -> np.ndarray:
-        weights, activations = self.get_cam_weights(input_tensor, target_layer, targets, activations, grads)
-
-        # 2D conv
-        if len(activations.shape) == 4:
-            weighted_activations = weights[:, :, None, None] * activations
-
-        # 3D conv
-        elif len(activations.shape) == 5:
-            weighted_activations = weights[:, :, None, None, None] * activations
-        else:
-            raise ValueError(f"Invalid activation shape. Get {len(activations.shape)}.")
-        
-        # weighted_activations = np.maximum(weighted_activations, 0)
-        # weighted_activations = np.abs(weighted_activations)
-        if eigen_smooth:
-            cam = get_2d_projection(weighted_activations)
-        else:
-            cam = weighted_activations.sum(axis=1)
-        return cam
-
-    def compute_cam_per_layer(
-        self, input_tensor: torch.Tensor, targets: List[torch.nn.Module], eigen_smooth: bool
-    ) -> np.ndarray:
-        activations_list = [a for a in self.activations_and_grads.original_activations]
-        grads_list = [g for g in self.activations_and_grads.original_gradients]
-        target_size = self.get_target_width_height(input_tensor)
-
-        cam_per_target_layer = []
-        # Loop over the saliency image from every layer
-        for i in range(len(self.target_layers)):
-            target_layer = self.target_layers[i]
-            layer_activations = None
-            layer_grads = None
-            if i < len(activations_list):
-                layer_activations = activations_list[i]
-            if i < len(grads_list):
-                layer_grads = grads_list[i]
-
-            cam = self.get_cam_image(input_tensor, target_layer, targets, layer_activations, layer_grads, eigen_smooth)
-            cam = np.maximum(cam, 0)
-            scaled = scale_cam_image(cam, target_size)
-            cam_per_target_layer.append(scaled[:, None, :])
-
-        return cam_per_target_layer

pytorch_grad_cam/shapley_cam.py

@@ -25,7 +25,7 @@ def __call__(self, model_output):
 
 class ClassifierOutputReST:
     """
-    Using both pre-softmax and post-softmax, propoesed in https://arxiv.org/abs/2501.06261
+    Using both pre-softmax and post-softmax, proposed in https://arxiv.org/abs/2501.06261
     """
     def __init__(self, category):
         self.category = category
@@ -36,7 +36,7 @@ def __call__(self, model_output):
             return model_output[0][self.category] - torch.nn.functional.cross_entropy(model_output, target)
         else:
             target = torch.tensor([self.category] * model_output.shape[0], device=model_output.device)
-            return model_output[:,self.category]- torch.nn.functional.cross_entropy(model_output, target)
+            return model_output[:,self.category] - torch.nn.functional.cross_entropy(model_output, target)
 
 
 class BinaryClassifierOutputTarget: