Pytorch class activation maps issue

veloc1rapt0r: <p>Hey everyone,</p>
<p>I’m encountering an issue while plotting class activation maps (CAMs) for a PyTorch neural network model. Specifically, I’m getting the following error:</p>
<pre><code class="lang-auto">RuntimeError: cannot register a hook on a tensor that doesn't require gradient.
</code></pre>
<p>I hope someone can help me clarify why this error occurs and how to resolve it <img src="https://discuss.pytorch.org/images/emoji/apple/blush.png?v=12" title=":blush:" class="emoji" alt=":blush:" loading="lazy" width="20" height="20">. The <code>plot_cam_model_list_images()</code> function is utilized within the training loop for supervised classification to visualize CAMs on the test set.</p>
<p>To replicate the error, I’ve created a test file and copied the code to Google Colab: <a href="https://colab.research.google.com/drive/1UU0r5LTvm6xZYs_xZH8hHL_gGXmiiwHW?usp=sharing" rel="noopener nofollow ugc">link to Colab notebook</a>.</p>
<p>Here’s a simplified version of the code:</p>
<pre data-code-wrap="python"><code class="lang-python">import matplotlib.pyplot as plt
import torch
import torchcam.methods as methods
from torchcam.utils import overlay_mask
from torchvision.transforms import ToPILImage
from torchvision.transforms.functional import to_pil_image
import torchvision.models as models
import torch.nn as nn

# Define ResNet18 model
class ResNet18(torch.nn.Module):
    def __init__(self, input_channels=1, output_channels=2):
        super().__init__()
        self.resnet = models.resnet18(pretrained=False)
        self.resnet.conv1 = nn.Conv2d(input_channels, 64, kernel_size=(
            7, 7), stride=(2, 2), padding=(3, 3), bias=False)
        self.resnet.fc = nn.Linear(
            in_features=512, out_features=output_channels, bias=True)

    def forward(self, x):
        return self.resnet(x)

def plot_cam_model_list_images(model, list_images):
    """
    Plot images with their corresponding CAM (Class Activation Mapping) overlays for a given model and list of images.
    Plots as a 2*number of images subplot with the first row containing the images and the second row containing the CAM overlays.
    Args:
        model (torch.nn.Module): The neural network model.
        list_images (list): List of images to visualize.
    Returns:
        None
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    input_shape = list_images[0].shape

    for param in model.parameters():
        param.requires_grad_(True)

    # Choose CAM method
    cam = methods.GradCAMpp(model=model, input_shape=input_shape)

    fig, ax = plt.subplots(nrows=2, ncols=len(list_images))
    ax = ax.flatten()
    plot_imgs_list = []
    plot_cams_list = []

    for i, image in enumerate(list_images):
        image.requires_grad = True
        image = image.to(device)

        # Forward pass
        out = model(image.unsqueeze(0))

        # Create a CAM extractor
        activation_map = cam(out.squeeze(0).argmax().item(), out)
        input_tensor_rgb = torch.cat([image, image, image], dim=0)

        # Move tensors to CPU
        input_tensor_rgb = input_tensor_rgb.cpu()

        # Overlay CAM on the input image
        cam_plot = overlay_mask(to_pil_image(input_tensor_rgb), to_pil_image(activation_map[0].squeeze(0), mode='F'),
                                alpha=0.5)

        plot_imgs_list.append(ToPILImage()(image))
        plot_cams_list.append(cam_plot)

    reordered_list_img_plot = plot_imgs_list + plot_cams_list
    for i, img in enumerate(reordered_list_img_plot):
        ax[i].imshow(img)
        ax[i].axis(False)
    plt.axis('off')
    plt.show()

# Usage example
model = ResNet18(1, 2)

for param in model.parameters():
    param.requires_grad = True

model.eval()

list_images1 = [torch.rand((1, 512, 512)) for i in range(4)]
plot_cam_model_list_images(model, list_images1)

list_images2 = [torch.rand((1, 512, 512)) for i in range(4)]
plot_cam_model_list_images(model, list_images2)
</code></pre>
<p>Upon running the <code>plot_cam_model_list_images()</code> function for the second time, I encounter the aforementioned <code>RuntimeError</code>. The traceback leads to the definition of CAM with the error being triggered during the registration of a hook.<br>
Full traceback:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "c:\Users\PythonWorkspace\VSCProjects\ml-incubator\source\deep_learning\main_scripts\test_CAM_script.py", line 93, in &lt;module&gt;
    plot_cam_model_list_images(
  File "c:\Users\PythonWorkspace\VSCProjects\ml-incubator\source\deep_learning\main_scripts\test_CAM_script.py", line 44, in plot_cam_model_list_images
    cam = methods.GradCAMpp(model=model, input_shape=input_shape)
  File "C:\python\python39enviroments\torch\lib\site-packages\torchcam\methods\gradient.py", line 33, in __init__
    super().__init__(model, target_layer, input_shape, **kwargs)
  File "C:\python\python39enviroments\torch\lib\site-packages\torchcam\methods\core.py", line 54, in __init__
    target_name = locate_candidate_layer(model, input_shape)
  File "C:\python\python39enviroments\torch\lib\site-packages\torchcam\methods\_utils.py", line 42, in locate_candidate_layer
    _ = mod(torch.zeros((1, *input_shape), device=next(mod.parameters()).data.device))
  File "C:\python\python39enviroments\torch\lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\python\python39enviroments\torch\lib\site-packages\torch\nn\modules\module.py", line 1561, in _call_impl
    result = forward_call(*args, **kwargs)
  File "c:\Users\\PythonWorkspace\VSCProjects\ml-incubator\source\deep_learning\main_scripts\test_CAM_script.py", line 22, in forward
    return self.resnet(x)
  File "C:\python\python39enviroments\torch\lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\python\python39enviroments\torch\lib\site-packages\torch\nn\modules\module.py", line 1561, in _call_impl
    result = forward_call(*args, **kwargs)
  File "C:\python\python39enviroments\torch\lib\site-packages\torchvision\models\resnet.py", line 285, in forward
    return self._forward_impl(x)
  File "C:\python\python39enviroments\torch\lib\site-packages\torchvision\models\resnet.py", line 276, in _forward_impl
    x = self.layer4(x)
  File "C:\python\python39enviroments\torch\lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\python\python39enviroments\torch\lib\site-packages\torch\nn\modules\module.py", line 1574, in _call_impl
    hook_result = hook(self, args, result)
  File "C:\python\python39enviroments\torch\lib\site-packages\torchcam\methods\gradient.py", line 49, in _hook_g
    self.hook_handles.append(output.register_hook(partial(self._store_grad, idx=idx)))
  File "C:\python\python39enviroments\torch\lib\site-packages\torch\_tensor.py", line 562, in register_hook
    raise RuntimeError(
RuntimeError: cannot register a hook on a tensor that doesn't require gradient
</code></pre>
<p>This issue appears similar to <a href="https://github.com/frgfm/torch-cam/discussions/132" rel="noopener nofollow ugc">this discussion</a>, but in my case, the model weights are not frozen.My model is a Pytorch ResNet18 changed for a given input and output channel dimensions (grayscale input and 2 classes at the end). In the real life main train function I would load the best checkpoint of the model before CAMs plotting.</p>
<p>Any assistance on resolving this issue would be greatly appreciated.</p>
<p>Thank you!</p>
<p>Best regards,</p>

veloc1rapt0r: <p>Hey everyone,<br>
I have fixed this issue by defining a single instance of CAM extractor for one script. Redefining the CAM messes up the hooks. I can’t find the thread I took the solution from, was somewhere in the Github discussions <a href="https://github.com/frgfm/torch-cam/discussions" class="inline-onebox" rel="noopener nofollow ugc">frgfm/torch-cam · Discussions · GitHub</a>. I have also updated the Colab notebook if somebody will face the same issue:)<br>
The issue is closed from my side.<br>
Solution:</p>
<pre><code>def plot_cam_model_list_images(model, cam, list_images):
"""
Plot images with their corresponding CAM (Class Activation Mapping) overlays for a given model and list of images.
Plots as a 2*number of images subplot with first row containing the images and the second row containing the CAM overlays.

Args:
    model (torch.nn.Module): The neural network model.
    list_images (list): List of images to visualize.
Returns:
    None
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

input_shape = list_images[0].shape

for param in model.parameters():
    param.requires_grad_(True)

fig, ax = plt.subplots(nrows=2, ncols=len(list_images))
ax = ax.flatten()
plot_imgs_list = []
plot_cams_list = []

for i, image in enumerate(list_images):
    image.requires_grad = True
    image = image.to(device)

    # Forward pass
    out = model(image.unsqueeze(0))

    # Create a CAM extractor
    activation_map = cam(out.squeeze(0).argmax().item(), out)
    input_tensor_rgb = torch.cat([image, image, image], dim=0)

    # Move tensors to CPU
    input_tensor_rgb = input_tensor_rgb.cpu()

    # Overlay CAM on the input image
    cam_plot = overlay_mask(to_pil_image(input_tensor_rgb), to_pil_image(activation_map[0].squeeze(0), mode='F'),
                            alpha=0.5)

    plot_imgs_list.append(ToPILImage()(image))
    plot_cams_list.append(cam_plot)

   reordered_list_img_plot = plot_imgs_list + plot_cams_list
   for i, img in enumerate(reordered_list_img_plot):
       ax[i].imshow(img)
       ax[i].axis(False)
   plt.axis('off')
   plt.show()


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ResNet18(1, 2)
model = model.to(device)

for param in model.parameters():
    param.requires_grad = True

model.eval()
input_shape = (1, 512, 512)
cam = methods.GradCAMpp(model=model, input_shape=input_shape)

list_images1 = [torch.rand(input_shape) for i in range(4)]
plot_cam_model_list_images(
    model, cam, list_images1)

list_images2 = [torch.rand(input_shape) for i in range(4)]
plot_cam_model_list_images(
    model, cam, list_images2)
</code></pre>