feat: Added automatic layer name resolution (#32)

* test: Renamed testers

* feat: Added layer resolution utils

* feat: Added layer resolution to CAM

* test: Added unittests for utils and simplified existing ones

* docs: Updated README

* refactor: Reflected changes on CAM interface

* feat: Updated visualization example script

* style: Fixed lint

* test: Fixed unittests

* feat: Forced eval switch when possible

README.md

@@ -15,7 +15,6 @@ Simple way to leverage the class-specific activation of convolutional layers in
  * [Prerequisites](#prerequisites)
  * [Installation](#installation)
 * [Usage](#usage)
-* [Technical Roadmap](#technical-roadmap)
 * [Documentation](#documentation)
 * [Contributing](#contributing)
 * [Credits](#credits)
@@ -58,20 +57,6 @@ python scripts/cam_example.py --model resnet50 --class-idx 232
 
 
 
-
-
-## Technical roadmap
-
-The project is currently under development, here are the objectives for the next releases:
-
-- [x] Parallel CAMs: enable batch processing.
-- [x] Benchmark: compare class activation map computations for different architectures.
-- [ ] Signature improvement: retrieve automatically the specific required layer names.
-- [ ] Refined RPN: create a region proposal network using CAM.
-- [ ] Task transfer: turn a well-trained classifier into an object detector.
-
-
-
 ## Documentation
 
 The full package documentation is available [here](https://frgfm.github.io/torch-cam/) for detailed specifications. The documentation was built with [Sphinx](sphinx-doc.org) using a [theme](github.com/readthedocs/sphinx_rtd_theme) provided by [Read the Docs](readthedocs.org).

scripts/cam_example.py

@@ -1,10 +1,10 @@
 #!usr/bin/python
-# -*- coding: utf-8 -*-
 
 """
 CAM visualization
 """
 
+import math
 import argparse
 from io import BytesIO
 
@@ -18,18 +18,18 @@
 from torchcam.cams import CAM, GradCAM, GradCAMpp, SmoothGradCAMpp, ScoreCAM, SSCAM, ISCAM
 from torchcam.utils import overlay_mask
 
-VGG_CONFIG = {_vgg: dict(input_layer='features', conv_layer='features')
+VGG_CONFIG = {_vgg: dict(conv_layer='features')
  for _vgg in models.vgg.__dict__.keys()}
 
-RESNET_CONFIG = {_resnet: dict(input_layer='conv1', conv_layer='layer4', fc_layer='fc')
+RESNET_CONFIG = {_resnet: dict(conv_layer='layer4', fc_layer='fc')
  for _resnet in models.resnet.__dict__.keys()}
 
-DENSENET_CONFIG = {_densenet: dict(input_layer='features', conv_layer='features', fc_layer='classifier')
+DENSENET_CONFIG = {_densenet: dict(conv_layer='features', fc_layer='classifier')
  for _densenet in models.densenet.__dict__.keys()}
 
 MODEL_CONFIG = {
  **VGG_CONFIG, **RESNET_CONFIG, **DENSENET_CONFIG,
- 'mobilenet_v2': dict(input_layer='features', conv_layer='features')
+ 'mobilenet_v2': dict(conv_layer='features')
 }
 
 
@@ -43,7 +43,6 @@ def main(args):
  # Pretrained imagenet model
  model = models.__dict__[args.model](pretrained=True).eval().to(device=device)
  conv_layer = MODEL_CONFIG[args.model]['conv_layer']
- input_layer = MODEL_CONFIG[args.model]['input_layer']
  fc_layer = MODEL_CONFIG[args.model]['fc_layer']
 
  # Image
@@ -57,15 +56,17 @@ def main(args):
 
  # Hook the corresponding layer in the model
  cam_extractors = [CAM(model, conv_layer, fc_layer), GradCAM(model, conv_layer),
- GradCAMpp(model, conv_layer), SmoothGradCAMpp(model, conv_layer, input_layer),
- ScoreCAM(model, conv_layer, input_layer), SSCAM(model, conv_layer, input_layer),
- ISCAM(model, conv_layer, input_layer)]
+ GradCAMpp(model, conv_layer), SmoothGradCAMpp(model, conv_layer),
+ ScoreCAM(model, conv_layer), SSCAM(model, conv_layer),
+ ISCAM(model, conv_layer)]
 
  # Don't trigger all hooks
  for extractor in cam_extractors:
  extractor._hooks_enabled = False
 
- fig, axes = plt.subplots(1, len(cam_extractors), figsize=(7, 2))
+ num_rows = 2
+ num_cols = math.ceil(len(cam_extractors) / num_rows)
+ _, axes = plt.subplots(num_rows, num_cols, figsize=(6, 4))
  for idx, extractor in enumerate(cam_extractors):
  extractor._hooks_enabled = True
  model.zero_grad()
@@ -76,6 +77,7 @@ def main(args):
 
  # Use the hooked data to compute activation map
  activation_map = extractor(class_idx, scores).cpu()
+
  # Clean data
  extractor.clear_hooks()
  extractor._hooks_enabled = False
@@ -85,9 +87,13 @@ def main(args):
  # Plot the result
  result = overlay_mask(pil_img, heatmap)
 
- axes[idx].imshow(result)
- axes[idx].axis('off')
- axes[idx].set_title(extractor.__class__.__name__, size=8)
+ axes[idx // num_cols][idx % num_cols].imshow(result)
+ axes[idx // num_cols][idx % num_cols].set_title(extractor.__class__.__name__, size=8)
+
+ # Clear axes
+ for row in axes:
+ for ax in row:
+ ax.axis('off')
 
  plt.tight_layout()
  if args.savefig:

test/test_cams.py

@@ -4,6 +4,7 @@
 import requests
 import torch
 from PIL import Image
+from torch import nn
 from torchvision.models import mobilenet_v2, resnet18
 from torchvision.transforms.functional import normalize, resize, to_tensor
 
@@ -20,7 +21,7 @@ def _forward(model, input_tensor):
  return scores
 
 
-class Tester(unittest.TestCase):
+class CAMCoreTester(unittest.TestCase):
  def _verify_cam(self, cam):
  # Simple verifications
  self.assertIsInstance(cam, torch.Tensor)
@@ -67,76 +68,91 @@ def _test_extractor(self, extractor, model):
 
  def _test_cam(self, name):
  # Get a pretrained model
- model = resnet18(pretrained=False).eval()
- conv_layer = 'layer4'
- input_layer = 'conv1'
- fc_layer = 'fc'
-
- # Hook the corresponding layer in the model
- extractor = cams.__dict__[name](model, conv_layer, fc_layer if name == 'CAM' else input_layer)
-
- self._test_extractor(extractor, model)
-
- def _test_cam_arbitrary_layer(self, name):
-
  model = resnet18(pretrained=False).eval()
  conv_layer = 'layer4.1.relu'
- input_layer = 'conv1'
- fc_layer = 'fc'
 
  # Hook the corresponding layer in the model
- extractor = cams.__dict__[name](model, conv_layer, fc_layer if name == 'CAM' else input_layer)
+ extractor = cams.__dict__[name](model, conv_layer)
 
- self._test_extractor(extractor, model)
+ with torch.no_grad():
+ self._test_extractor(extractor, model)
 
  def _test_gradcam(self, name):
 
  # Get a pretrained model
  model = mobilenet_v2(pretrained=False)
- conv_layer = 'features'
+ conv_layer = 'features.17.conv.3'
 
  # Hook the corresponding layer in the model
  extractor = cams.__dict__[name](model, conv_layer)
 
  self._test_extractor(extractor, model)
 
- def _test_gradcam_arbitrary_layer(self, name):
+ def test_smooth_gradcampp(self):
 
- model = mobilenet_v2(pretrained=False)
- conv_layer = 'features.17.conv.3'
+ # Get a pretrained model
+ model = mobilenet_v2(pretrained=False).eval()
 
  # Hook the corresponding layer in the model
- extractor = cams.__dict__[name](model, conv_layer)
+ extractor = cams.SmoothGradCAMpp(model)
 
  self._test_extractor(extractor, model)
 
- def test_smooth_gradcampp(self):
 
- # Get a pretrained model
- model = mobilenet_v2(pretrained=False)
- conv_layer = 'features'
- input_layer = 'features'
+class CAMUtilsTester(unittest.TestCase):
 
-  # Hook the corresponding layer in the model
-  extractor = cams.SmoothGradCAMpp(model, conv_layer, input_layer)
+ @staticmethod
+ def _get_custom_module():
 
- self._test_extractor(extractor, model)
+ mod = nn.Sequential(
+ nn.Sequential(
+ nn.Conv2d(3, 8, 3, 1),
+ nn.ReLU(),
+ nn.Conv2d(8, 16, 3, 1),
+ nn.ReLU(),
+ nn.AdaptiveAvgPool2d((1, 1))
+ ),
+ nn.Flatten(1),
+ nn.Linear(16, 1)
+ )
+ return mod
+
+ def test_locate_candidate_layer(self):
+
+ # ResNet-18
+ mod = resnet18().eval()
+ self.assertEqual(cams.utils.locate_candidate_layer(mod), 'layer4')
+
+ # Custom model
+ mod = self._get_custom_module()
+
+ self.assertEqual(cams.utils.locate_candidate_layer(mod), '0.3')
+ # Check that the model is switched back to its origin mode afterwards
+ self.assertTrue(mod.training)
+
+ def test_locate_linear_layer(self):
+
+ # ResNet-18
+ mod = resnet18().eval()
+ self.assertEqual(cams.utils.locate_linear_layer(mod), 'fc')
+
+ # Custom model
+ mod = self._get_custom_module()
+ self.assertEqual(cams.utils.locate_linear_layer(mod), '2')
 
 
 for cam_extractor in ['CAM', 'ScoreCAM', 'SSCAM', 'ISCAM']:
  def do_test(self, cam_extractor=cam_extractor):
  self._test_cam(cam_extractor)
- self._test_cam_arbitrary_layer(cam_extractor)
 
- setattr(Tester, "test_" + cam_extractor.lower(), do_test)
+ setattr(CAMCoreTester, "test_" + cam_extractor.lower(), do_test)
 
 
 for cam_extractor in ['GradCAM', 'GradCAMpp']:
  def do_test(self, cam_extractor=cam_extractor):
  self._test_gradcam(cam_extractor)
- self._test_gradcam_arbitrary_layer(cam_extractor)
 
- setattr(Tester, "test_" + cam_extractor.lower(), do_test)
+ setattr(CAMCoreTester, "test_" + cam_extractor.lower(), do_test)
 
 
 if __name__ == '__main__':

test/test_utils.py

@@ -6,7 +6,7 @@
 from torchcam import utils
 
 
-class Tester(unittest.TestCase):
+class UtilsTester(unittest.TestCase):
  def test_overlay_mask(self):
 
  img = Image.fromarray(np.zeros((4, 4, 3)).astype(np.uint8))

torchcam/cams/__init__.py

@@ -1,2 +1,3 @@
 from .cam import *
 from .gradcam import *
+from .utils import *