ShapleyCAM

Weighting the activation maps using Gradient and Hessian-Vector Product.

Author: cai2-huaiguang

README.md

@@ -47,6 +47,7 @@ The aim is also to serve as a benchmark of algorithms and metrics for research o
 | Deep Feature Factorizations           | Non Negative Matrix Factorization on the 2D activations                                                   |
 | KPCA-CAM            | Like EigenCAM but with Kernel PCA instead of PCA                                                                            |            
 | FEM                 | A gradient free method that binarizes activations by an activation > mean + k * std rule.                                   |
+| ShapleyCAM          | Weighting the activation maps using Gradient and Hessian-Vector Product.|
 ## Visual Examples
 
 | What makes the network think the image label is 'pug, pug-dog' | What makes the network think the image label is 'tabby, tabby cat' | Combining Grad-CAM with Guided Backpropagation for the 'pug, pug-dog' class |

cam.py

@@ -7,7 +7,7 @@
 from pytorch_grad_cam import (
     GradCAM, FEM, HiResCAM, ScoreCAM, GradCAMPlusPlus,
     AblationCAM, XGradCAM, EigenCAM, EigenGradCAM,
-    LayerCAM, FullGrad, GradCAMElementWise, KPCA_CAM
+    LayerCAM, FullGrad, GradCAMElementWise, KPCA_CAM, ShapleyCAM
 )
 from pytorch_grad_cam import GuidedBackpropReLUModel
 from pytorch_grad_cam.utils.image import (
@@ -37,7 +37,7 @@ def get_args():
                             'gradcam', 'fem', 'hirescam', 'gradcam++',
                             'scorecam', 'xgradcam', 'ablationcam',
                             'eigencam', 'eigengradcam', 'layercam',
-                            'fullgrad', 'gradcamelementwise', 'kpcacam'
+                            'fullgrad', 'gradcamelementwise', 'kpcacam', 'shapleycam'
                         ],
                         help='CAM method')
 
@@ -75,7 +75,8 @@ def get_args():
         "fullgrad": FullGrad,
         "fem": FEM,
         "gradcamelementwise": GradCAMElementWise,
-        'kpcacam': KPCA_CAM
+        'kpcacam': KPCA_CAM,
+        'shapleycam': ShapleyCAM
     }
 
     if args.device=='hpu':

pytorch_grad_cam/__init__.py

@@ -1,4 +1,5 @@
 from pytorch_grad_cam.grad_cam import GradCAM
+from pytorch_grad_cam.shapley_cam import ShapleyCAM
 from pytorch_grad_cam.fem import FEM
 from pytorch_grad_cam.hirescam import HiResCAM
 from pytorch_grad_cam.grad_cam_elementwise import GradCAMElementWise

pytorch_grad_cam/activations_and_gradients_no_detach.py

@@ -0,0 +1,53 @@
+class ActivationsAndGradients_no_detach:
+    """ Class for extracting activations and
+    registering gradients from targetted intermediate layers """
+
+    def __init__(self, model, target_layers, reshape_transform):
+        self.model = model
+        # self.gradients = []
+        # self.activations = []
+        self.original_gradients = []
+        self.original_activations = []
+        self.reshape_transform = reshape_transform
+        self.handles = []
+        for target_layer in target_layers:
+            self.handles.append(
+                target_layer.register_forward_hook(self.save_activation))
+            # Because of https://github.com/pytorch/pytorch/issues/61519,
+            # we don't use backward hook to record gradients.
+            self.handles.append(
+                target_layer.register_forward_hook(self.save_gradient))
+
+    def save_activation(self, module, input, output):
+        activation = output
+
+        self.original_activations.append(activation)
+        # if self.reshape_transform is not None:
+        #     activation = self.reshape_transform(activation)
+        # # self.activations.append(activation.cpu().detach())
+        # self.activations.append(activation)
+
+    def save_gradient(self, module, input, output):
+        if not hasattr(output, "requires_grad") or not output.requires_grad:
+            # You can only register hooks on tensor requires grad.
+            return
+
+        # Gradients are computed in reverse order
+        def _store_grad(grad):
+            self.original_gradients = [grad] + self.original_gradients
+            # if self.reshape_transform is not None:
+            #     grad = self.reshape_transform(grad)
+            # self.gradients = [grad] + self.gradients
+
+        output.register_hook(_store_grad)
+
+    def __call__(self, x):
+        # self.gradients = []
+        # self.activations = []
+        self.original_gradients = []
+        self.original_activations = []
+        return self.model(x)
+
+    def release(self):
+        for handle in self.handles:
+            handle.remove()

pytorch_grad_cam/shapley_cam.py

@@ -0,0 +1,159 @@
+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
+
+"""
+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.
+"""
+class ShapleyCAM(BaseCAM):
+    def __init__(self, model, target_layers,
+                 reshape_transform=None):
+        super(
+            ShapleyCAM,
+            self).__init__(
+            model,
+            target_layers,
+            reshape_transform)
+
+        self.activations_and_grads = ActivationsAndGradients_no_detach(self.model, target_layers, reshape_transform)
+
+    def forward(
+        self, input_tensor: torch.Tensor, targets: List[torch.nn.Module], eigen_smooth: bool = False
+    ) -> np.ndarray:
+        input_tensor = input_tensor.to(self.device)
+
+        input_tensor = torch.autograd.Variable(input_tensor, requires_grad=True)
+
+        self.outputs = outputs = self.activations_and_grads(input_tensor)
+
+        if targets is None:
+            target_categories = np.argmax(outputs.cpu().data.numpy(), axis=-1)
+            targets = [ClassifierOutputTarget(category) for category in target_categories]
+
+        if self.uses_gradients:
+            self.model.zero_grad()
+            loss = sum([target(output) for target, output in zip(targets, outputs)])
+            torch.autograd.grad(loss, input_tensor,  retain_graph = True, create_graph = True)
+
+        # In most of the saliency attribution papers, the saliency is
+        # computed with a single target layer.
+        # Commonly it is the last convolutional layer.
+        # Here we support passing a list with multiple target layers.
+        # It will compute the saliency image for every image,
+        # and then aggregate them (with a default mean aggregation).
+        # This gives you more flexibility in case you just want to
+        # use all conv layers for example, all Batchnorm layers,
+        # or something else.
+        cam_per_layer = self.compute_cam_per_layer(input_tensor, targets, eigen_smooth)
+        return self.aggregate_multi_layers(cam_per_layer)
+
+
+    def get_cam_weights(self,
+                        input_tensor,
+                        target_layer,
+                        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]
+        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)
+
+        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()
+
+
+        # 2D image
+        if len(activations.shape) == 4:
+            weight = np.mean(weight, axis=(2, 3))
+            return weight, activations
+        
+        # 3D image
+        elif len(activations.shape) == 5:
+            weight = np.mean(weight, axis=(2, 3, 4))
+            return weight, activations
+        
+        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/utils/model_targets.py

@@ -23,6 +23,22 @@ def __call__(self, model_output):
         return torch.softmax(model_output, dim=-1)[:, self.category]
 
 
+class ClassifierOutputReST:
+    """
+    Using both pre-softmax and post-softmax, propoesed in https://arxiv.org/abs/2501.06261
+    """
+    def __init__(self, category):
+        self.category = category
+    def __call__(self, model_output): 
+        if len(model_output.shape) == 1:
+            target = torch.tensor([self.category], device=model_output.device)
+            model_output = model_output.unsqueeze(0)
+            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)
+
+
 class BinaryClassifierOutputTarget:
     def __init__(self, category):
         self.category = category