GradCAM not working with cryptic error message

[mattwarkentin]

Describe the bug
I am trying to use the GradCAM on my model but I can't seem to get it working and the error isn't really helping me figure out the issue. Hoping someone here could provide some insight.

For context, this model is based on a resnet50 backbone, so the input is RGB and size 224 x 224.

To Reproduce

from monai.visualize.class_activation_maps import GradCAM
from py.main import MyModule, MyDataModule
from glob import glob

best_model = glob("lightning_logs/chkpt/*.ckpt")[0]

model = MyModule.load_from_checkpoint(best_model)

datamodule = MyDataModule(batch_size=1)

batch = next(iter(datamodule.val_dataloader()))
x, y = model.prepare_batch(batch)

cam = GradCAM(model, target_layers="model.layer4.2.relu")

model(x) ## this works

result = cam(x, class_idx=0)

Traceback (most recent call last):
 File "<stdin>", line 1, in <module>
 File "/Users/Matt/.virtualenvs/torch/lib/python3.8/site-packages/monai/visualize/class_activation_maps.py", line 371, in __call__
   acti_map = self.compute_map(x, class_idx=class_idx, retain_graph=retain_graph, layer_idx=layer_idx)
 File "/Users/Matt/.virtualenvs/torch/lib/python3.8/site-packages/monai/visualize/class_activation_maps.py", line 351, in compute_map
   _, acti, grad = self.nn_module(x, class_idx=class_idx, retain_graph=retain_graph)
 File "/Users/Matt/.virtualenvs/torch/lib/python3.8/site-packages/monai/visualize/class_activation_maps.py", line 135, in __call__
   acti = tuple(self.activations[layer] for layer in self.target_layers)
 File "/Users/Matt/.virtualenvs/torch/lib/python3.8/site-packages/monai/visualize/class_activation_maps.py", line 135, in <genexpr>
   acti = tuple(self.activations[layer] for layer in self.target_layers)
KeyError: 'layer4.2.relu’

for name, _ in model.named_modules(): print(name)

model
model.conv1
model.bn1
model.relu
model.maxpool
model.layer1
model.layer1.0
model.layer1.0.conv1
model.layer1.0.bn1
model.layer1.0.conv2
model.layer1.0.bn2
model.layer1.0.conv3
model.layer1.0.bn3
model.layer1.0.relu
model.layer1.0.downsample
model.layer1.0.downsample.0
model.layer1.0.downsample.1
model.layer1.1
model.layer1.1.conv1
model.layer1.1.bn1
model.layer1.1.conv2
model.layer1.1.bn2
model.layer1.1.conv3
model.layer1.1.bn3
model.layer1.1.relu
model.layer1.2
model.layer1.2.conv1
model.layer1.2.bn1
model.layer1.2.conv2
model.layer1.2.bn2
model.layer1.2.conv3
model.layer1.2.bn3
model.layer1.2.relu
model.layer2
model.layer2.0
model.layer2.0.conv1
model.layer2.0.bn1
model.layer2.0.conv2
model.layer2.0.bn2
model.layer2.0.conv3
model.layer2.0.bn3
model.layer2.0.relu
model.layer2.0.downsample
model.layer2.0.downsample.0
model.layer2.0.downsample.1
model.layer2.1
model.layer2.1.conv1
model.layer2.1.bn1
model.layer2.1.conv2
model.layer2.1.bn2
model.layer2.1.conv3
model.layer2.1.bn3
model.layer2.1.relu
model.layer2.2
model.layer2.2.conv1
model.layer2.2.bn1
model.layer2.2.conv2
model.layer2.2.bn2
model.layer2.2.conv3
model.layer2.2.bn3
model.layer2.2.relu
model.layer2.3
model.layer2.3.conv1
model.layer2.3.bn1
model.layer2.3.conv2
model.layer2.3.bn2
model.layer2.3.conv3
model.layer2.3.bn3
model.layer2.3.relu
model.layer3
model.layer3.0
model.layer3.0.conv1
model.layer3.0.bn1
model.layer3.0.conv2
model.layer3.0.bn2
model.layer3.0.conv3
model.layer3.0.bn3
model.layer3.0.relu
model.layer3.0.downsample
model.layer3.0.downsample.0
model.layer3.0.downsample.1
model.layer3.1
model.layer3.1.conv1
model.layer3.1.bn1
model.layer3.1.conv2
model.layer3.1.bn2
model.layer3.1.conv3
model.layer3.1.bn3
model.layer3.1.relu
model.layer3.2
model.layer3.2.conv1
model.layer3.2.bn1
model.layer3.2.conv2
model.layer3.2.bn2
model.layer3.2.conv3
model.layer3.2.bn3
model.layer3.2.relu
model.layer3.3
model.layer3.3.conv1
model.layer3.3.bn1
model.layer3.3.conv2
model.layer3.3.bn2
model.layer3.3.conv3
model.layer3.3.bn3
model.layer3.3.relu
model.layer3.4
model.layer3.4.conv1
model.layer3.4.bn1
model.layer3.4.conv2
model.layer3.4.bn2
model.layer3.4.conv3
model.layer3.4.bn3
model.layer3.4.relu
model.layer3.5
model.layer3.5.conv1
model.layer3.5.bn1
model.layer3.5.conv2
model.layer3.5.bn2
model.layer3.5.conv3
model.layer3.5.bn3
model.layer3.5.relu
model.layer4
model.layer4.0
model.layer4.0.conv1
model.layer4.0.bn1
model.layer4.0.conv2
model.layer4.0.bn2
model.layer4.0.conv3
model.layer4.0.bn3
model.layer4.0.relu
model.layer4.0.downsample
model.layer4.0.downsample.0
model.layer4.0.downsample.1
model.layer4.1
model.layer4.1.conv1
model.layer4.1.bn1
model.layer4.1.conv2
model.layer4.1.bn2
model.layer4.1.conv3
model.layer4.1.bn3
model.layer4.1.relu
model.layer4.2
model.layer4.2.conv1
model.layer4.2.bn1
model.layer4.2.conv2
model.layer4.2.bn2
model.layer4.2.conv3
model.layer4.2.bn3
model.layer4.2.relu
model.avgpool
model.fc
model.fc.0
model.fc.1

Expected behavior
Expected that a tensor is returned of the same dimension as x.

Environment

================================
Printing MONAI config...
================================
MONAI version: 0.6.0
Numpy version: 1.19.5
Pytorch version: 1.9.1
MONAI flags: HAS_EXT = False, USE_COMPILED = False
MONAI rev id: 0ad9e73639e30f4f1af5a1f4a45da9cb09930179

Optional dependencies:
Pytorch Ignite version: NOT INSTALLED or UNKNOWN VERSION.
Nibabel version: 3.2.1
scikit-image version: NOT INSTALLED or UNKNOWN VERSION.
Pillow version: 8.3.2
Tensorboard version: 2.6.0
gdown version: NOT INSTALLED or UNKNOWN VERSION.
TorchVision version: 0.10.1
ITK version: 5.2.1
tqdm version: 4.62.3
lmdb version: NOT INSTALLED or UNKNOWN VERSION.
psutil version: NOT INSTALLED or UNKNOWN VERSION.
pandas version: 1.2.5
einops version: NOT INSTALLED or UNKNOWN VERSION.

[wyli]

Strange here the param is target_layers="model.layer4.2.relu" but the error is about KeyError: 'layer4.2.relu’. It's not possible to reproduce the issue without MyModule, MyDataModule, could you make a self-contained example? I can then look into it.

[mattwarkentin]

@wyli Here is a minimal reproducible example:

from monai.visualize.class_activation_maps import GradCAM
import pytorch_lightning as pl
import torchvision as tv
import torch

class MyModule(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = tv.models.resnet50(pretrained=True)

        ## Freeze backbone
        for param in self.model.parameters():
            param.requires_grad = False
        
        ## Change head to single class
        self.model.fc = torch.nn.Linear(self.model.fc.in_features, 1)
    
    def forward(self, x):
        return self.model(x)

x = torch.rand((1, 3, 224, 224))

model = MyModule()

cam = GradCAM(model, target_layers='model.layer4.2.conv3')

cam(x)

[mattwarkentin]

Seems like the issue is related to freezing the backbone as the GradCAM works when I remove the for loop. I do actually want the backbone frozen during training, though, so how should I handle this when wanting to use CAM?

Neither of these work...

with torch.enable_grad():
    cam(x)

with torch.set_grad_enabled(True):
    cam(x)

[rijobro]

Is it not because you have a model wrapped inside a model?

[mattwarkentin]

I don't think so. A PyTorch Lightning module is a nn_module so it should work as such.

The code below still fails.

model = tv.models.resnet50(pretrained=True)

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

model.fc = torch.nn.Linear(model.fc.in_features, 1)

cam = GradCAM(model, target_layers='layer4.2.conv3')

cam(x)

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/matt/.virtualenvs/torch/lib/python3.7/site-packages/monai/visualize/class_activation_maps.py", line 364, in __call__
    acti_map = self.compute_map(x, class_idx=class_idx, retain_graph=retain_graph, layer_idx=layer_idx)
  File "/Users/matt/.virtualenvs/torch/lib/python3.7/site-packages/monai/visualize/class_activation_maps.py", line 344, in compute_map
    _, acti, grad = self.nn_module(x, class_idx=class_idx, retain_graph=retain_graph)
  File "/Users/matt/.virtualenvs/torch/lib/python3.7/site-packages/monai/visualize/class_activation_maps.py", line 133, in __call__
    grad = tuple(self.gradients[layer] for layer in self.target_layers)
  File "/Users/matt/.virtualenvs/torch/lib/python3.7/site-packages/monai/visualize/class_activation_maps.py", line 133, in <genexpr>
    grad = tuple(self.gradients[layer] for layer in self.target_layers)
KeyError: 'layer4.2.conv3'

[wyli]

I tried a few things and this one seems to work:

from monai.visualize.class_activation_maps import GradCAM
import pytorch_lightning as pl
import torchvision as tv
import torch

class MyModule(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = tv.models.resnet50(pretrained=True)

        # # Freeze backbone
        for name, param in self.model.named_parameters():
            if "layer4" not in name:
                param.requires_grad = False

        # Change head to single class
        self.model.fc = torch.nn.Linear(self.model.fc.in_features, 1)

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

x = torch.rand((1, 3, 224, 224))

model = MyModule()

cam = GradCAM(model, target_layers='model.layer4')

print(cam(x).shape)

[mattwarkentin]

If the gradients can't be computed (and thus CAM can't be used) if requires_grad is False for the target layer(s), would it make sense for GradCAM et al. to either throw a warning/error if this isn't true during __init__, or automatically set requires_grad to true?

[wyli]

sounds good, I'll enhance the error handling.