files

pytorch_grad_cam/__init__.py

@@ -0,0 +1,20 @@
+from pytorch_grad_cam.grad_cam import GradCAM
+from pytorch_grad_cam.hirescam import HiResCAM
+from pytorch_grad_cam.grad_cam_elementwise import GradCAMElementWise
+from pytorch_grad_cam.ablation_layer import AblationLayer, AblationLayerVit, AblationLayerFasterRCNN
+from pytorch_grad_cam.ablation_cam import AblationCAM
+from pytorch_grad_cam.xgrad_cam import XGradCAM
+from pytorch_grad_cam.grad_cam_plusplus import GradCAMPlusPlus
+from pytorch_grad_cam.score_cam import ScoreCAM
+from pytorch_grad_cam.layer_cam import LayerCAM
+from pytorch_grad_cam.eigen_cam import EigenCAM
+from pytorch_grad_cam.eigen_grad_cam import EigenGradCAM
+from pytorch_grad_cam.random_cam import RandomCAM
+from pytorch_grad_cam.fullgrad_cam import FullGrad
+from pytorch_grad_cam.guided_backprop import GuidedBackpropReLUModel
+from pytorch_grad_cam.activations_and_gradients import ActivationsAndGradients
+from pytorch_grad_cam.feature_factorization.deep_feature_factorization import DeepFeatureFactorization, run_dff_on_image
+import pytorch_grad_cam.utils.model_targets
+import pytorch_grad_cam.utils.reshape_transforms
+import pytorch_grad_cam.metrics.cam_mult_image
+import pytorch_grad_cam.metrics.road

pytorch_grad_cam/ablation_cam.py

@@ -0,0 +1,148 @@
+import numpy as np
+import torch
+import tqdm
+from typing import Callable, List
+from pytorch_grad_cam.base_cam import BaseCAM
+from pytorch_grad_cam.utils.find_layers import replace_layer_recursive
+from pytorch_grad_cam.ablation_layer import AblationLayer
+
+
+""" Implementation of AblationCAM
+https://openaccess.thecvf.com/content_WACV_2020/papers/Desai_Ablation-CAM_Visual_Explanations_for_Deep_Convolutional_Network_via_Gradient-free_Localization_WACV_2020_paper.pdf
+
+Ablate individual activations, and then measure the drop in the target score.
+
+In the current implementation, the target layer activations is cached, so it won't be re-computed.
+However layers before it, if any, will not be cached.
+This means that if the target layer is a large block, for example model.featuers (in vgg), there will
+be a large save in run time.
+
+Since we have to go over many channels and ablate them, and every channel ablation requires a forward pass,
+it would be nice if we could avoid doing that for channels that won't contribute anwyay, making it much faster.
+The parameter ratio_channels_to_ablate controls how many channels should be ablated, using an experimental method
+(to be improved). The default 1.0 value means that all channels will be ablated.
+"""
+
+
+class AblationCAM(BaseCAM):
+    def __init__(self,
+                 model: torch.nn.Module,
+                 target_layers: List[torch.nn.Module],
+                 use_cuda: bool = False,
+                 reshape_transform: Callable = None,
+                 ablation_layer: torch.nn.Module = AblationLayer(),
+                 batch_size: int = 32,
+                 ratio_channels_to_ablate: float = 1.0) -> None:
+
+        super(AblationCAM, self).__init__(model,
+                                          target_layers,
+                                          use_cuda,
+                                          reshape_transform,
+                                          uses_gradients=False)
+        self.batch_size = batch_size
+        self.ablation_layer = ablation_layer
+        self.ratio_channels_to_ablate = ratio_channels_to_ablate
+
+    def save_activation(self, module, input, output) -> None:
+        """ Helper function to save the raw activations from the target layer """
+        self.activations = output
+
+    def assemble_ablation_scores(self,
+                                 new_scores: list,
+                                 original_score: float,
+                                 ablated_channels: np.ndarray,
+                                 number_of_channels: int) -> np.ndarray:
+        """ Take the value from the channels that were ablated,
+            and just set the original score for the channels that were skipped """
+
+        index = 0
+        result = []
+        sorted_indices = np.argsort(ablated_channels)
+        ablated_channels = ablated_channels[sorted_indices]
+        new_scores = np.float32(new_scores)[sorted_indices]
+
+        for i in range(number_of_channels):
+            if index < len(ablated_channels) and ablated_channels[index] == i:
+                weight = new_scores[index]
+                index = index + 1
+            else:
+                weight = original_score
+            result.append(weight)
+
+        return result
+
+    def get_cam_weights(self,
+                        input_tensor: torch.Tensor,
+                        target_layer: torch.nn.Module,
+                        targets: List[Callable],
+                        activations: torch.Tensor,
+                        grads: torch.Tensor) -> np.ndarray:
+
+        # Do a forward pass, compute the target scores, and cache the
+        # activations
+        handle = target_layer.register_forward_hook(self.save_activation)
+        with torch.no_grad():
+            outputs = self.model(input_tensor)
+            handle.remove()
+            original_scores = np.float32(
+                [target(output).cpu().item() for target, output in zip(targets, outputs)])
+
+        # Replace the layer with the ablation layer.
+        # When we finish, we will replace it back, so the original model is
+        # unchanged.
+        ablation_layer = self.ablation_layer
+        replace_layer_recursive(self.model, target_layer, ablation_layer)
+
+        number_of_channels = activations.shape[1]
+        weights = []
+        # This is a "gradient free" method, so we don't need gradients here.
+        with torch.no_grad():
+            # Loop over each of the batch images and ablate activations for it.
+            for batch_index, (target, tensor) in enumerate(
+                    zip(targets, input_tensor)):
+                new_scores = []
+                batch_tensor = tensor.repeat(self.batch_size, 1, 1, 1)
+
+                # Check which channels should be ablated. Normally this will be all channels,
+                # But we can also try to speed this up by using a low
+                # ratio_channels_to_ablate.
+                channels_to_ablate = ablation_layer.activations_to_be_ablated(
+                    activations[batch_index, :], self.ratio_channels_to_ablate)
+                number_channels_to_ablate = len(channels_to_ablate)
+
+                for i in tqdm.tqdm(
+                    range(
+                        0,
+                        number_channels_to_ablate,
+                        self.batch_size)):
+                    if i + self.batch_size > number_channels_to_ablate:
+                        batch_tensor = batch_tensor[:(
+                            number_channels_to_ablate - i)]
+
+                    # Change the state of the ablation layer so it ablates the next channels.
+                    # TBD: Move this into the ablation layer forward pass.
+                    ablation_layer.set_next_batch(
+                        input_batch_index=batch_index,
+                        activations=self.activations,
+                        num_channels_to_ablate=batch_tensor.size(0))
+                    score = [target(o).cpu().item()
+                             for o in self.model(batch_tensor)]
+                    new_scores.extend(score)
+                    ablation_layer.indices = ablation_layer.indices[batch_tensor.size(
+                        0):]
+
+                new_scores = self.assemble_ablation_scores(
+                    new_scores,
+                    original_scores[batch_index],
+                    channels_to_ablate,
+                    number_of_channels)
+                weights.extend(new_scores)
+
+        weights = np.float32(weights)
+        weights = weights.reshape(activations.shape[:2])
+        original_scores = original_scores[:, None]
+        weights = (original_scores - weights) / original_scores
+
+        # Replace the model back to the original state
+        replace_layer_recursive(self.model, ablation_layer, target_layer)
+        return weights

pytorch_grad_cam/ablation_cam_multilayer.py

@@ -0,0 +1,136 @@
+import cv2
+import numpy as np
+import torch
+import tqdm
+from pytorch_grad_cam.base_cam import BaseCAM
+
+
+class AblationLayer(torch.nn.Module):
+    def __init__(self, layer, reshape_transform, indices):
+        super(AblationLayer, self).__init__()
+
+        self.layer = layer
+        self.reshape_transform = reshape_transform
+        # The channels to zero out:
+        self.indices = indices
+
+    def forward(self, x):
+        self.__call__(x)
+
+    def __call__(self, x):
+        output = self.layer(x)
+
+        # Hack to work with ViT,
+        # Since the activation channels are last and not first like in CNNs
+        # Probably should remove it?
+        if self.reshape_transform is not None:
+            output = output.transpose(1, 2)
+
+        for i in range(output.size(0)):
+
+            # Commonly the minimum activation will be 0,
+            # And then it makes sense to zero it out.
+            # However depending on the architecture,
+            # If the values can be negative, we use very negative values
+            # to perform the ablation, deviating from the paper.
+            if torch.min(output) == 0:
+                output[i, self.indices[i], :] = 0
+            else:
+                ABLATION_VALUE = 1e5
+                output[i, self.indices[i], :] = torch.min(
+                    output) - ABLATION_VALUE
+
+        if self.reshape_transform is not None:
+            output = output.transpose(2, 1)
+
+        return output
+
+
+def replace_layer_recursive(model, old_layer, new_layer):
+    for name, layer in model._modules.items():
+        if layer == old_layer:
+            model._modules[name] = new_layer
+            return True
+        elif replace_layer_recursive(layer, old_layer, new_layer):
+            return True
+    return False
+
+
+class AblationCAM(BaseCAM):
+    def __init__(self, model, target_layers, use_cuda=False,
+                 reshape_transform=None):
+        super(AblationCAM, self).__init__(model, target_layers, use_cuda,
+                                          reshape_transform)
+
+        if len(target_layers) > 1:
+            print(
+                "Warning. You are usign Ablation CAM with more than 1 layers. "
+                "This is supported only if all layers have the same output shape")
+
+    def set_ablation_layers(self):
+        self.ablation_layers = []
+        for target_layer in self.target_layers:
+            ablation_layer = AblationLayer(target_layer,
+                                           self.reshape_transform, indices=[])
+            self.ablation_layers.append(ablation_layer)
+            replace_layer_recursive(self.model, target_layer, ablation_layer)
+
+    def unset_ablation_layers(self):
+        # replace the model back to the original state
+        for ablation_layer, target_layer in zip(
+                self.ablation_layers, self.target_layers):
+            replace_layer_recursive(self.model, ablation_layer, target_layer)
+
+    def set_ablation_layer_batch_indices(self, indices):
+        for ablation_layer in self.ablation_layers:
+            ablation_layer.indices = indices
+
+    def trim_ablation_layer_batch_indices(self, keep):
+        for ablation_layer in self.ablation_layers:
+            ablation_layer.indices = ablation_layer.indices[:keep]
+
+    def get_cam_weights(self,
+                        input_tensor,
+                        target_category,
+                        activations,
+                        grads):
+        with torch.no_grad():
+            outputs = self.model(input_tensor).cpu().numpy()
+            original_scores = []
+            for i in range(input_tensor.size(0)):
+                original_scores.append(outputs[i, target_category[i]])
+        original_scores = np.float32(original_scores)
+
+        self.set_ablation_layers()
+
+        if hasattr(self, "batch_size"):
+            BATCH_SIZE = self.batch_size
+        else:
+            BATCH_SIZE = 32
+
+        number_of_channels = activations.shape[1]
+        weights = []
+
+        with torch.no_grad():
+            # Iterate over the input batch
+            for tensor, category in zip(input_tensor, target_category):
+                batch_tensor = tensor.repeat(BATCH_SIZE, 1, 1, 1)
+                for i in tqdm.tqdm(range(0, number_of_channels, BATCH_SIZE)):
+                    self.set_ablation_layer_batch_indices(
+                        list(range(i, i + BATCH_SIZE)))
+
+                    if i + BATCH_SIZE > number_of_channels:
+                        keep = number_of_channels - i
+                        batch_tensor = batch_tensor[:keep]
+                        self.trim_ablation_layer_batch_indices(self, keep)
+                    score = self.model(batch_tensor)[:, category].cpu().numpy()
+                    weights.extend(score)
+
+        weights = np.float32(weights)
+        weights = weights.reshape(activations.shape[:2])
+        original_scores = original_scores[:, None]
+        weights = (original_scores - weights) / original_scores
+
+        # replace the model back to the original state
+        self.unset_ablation_layers()
+        return weights