docs: Fixed example docstring and unittests (#33)

* docs: Fixed example of SmoothGradCAMpp

* test: Fixed unittests

* test: Speeded up unittests

* test: Fixed base Cam unittest

* test: Optimized speed for Score CAM family

* test: Optimized testing speed for SSCAM and ISCAM

test/test_cams.py

@@ -11,16 +11,6 @@
 from torchcam import cams
 
 
-def _forward(model, input_tensor):
- if model.training:
- scores = model(input_tensor)
- else:
- with torch.no_grad():
- scores = model(input_tensor)
-
- return scores
-
-
 class CAMCoreTester(unittest.TestCase):
  def _verify_cam(self, cam):
  # Simple verifications
@@ -50,11 +40,11 @@ def _test_extractor(self, extractor, model):
  img_tensor = self._get_img_tensor()
 
  # Check that a batch of 2 cannot be accepted
- _ = _forward(model, torch.stack((img_tensor, img_tensor)))
+ _ = model(torch.stack((img_tensor, img_tensor)))
  self.assertRaises(ValueError, extractor, 0)
 
  # Correct forward
- scores = _forward(model, img_tensor.unsqueeze(0))
+ scores = model(img_tensor.unsqueeze(0))
 
  # Check incorrect class index
  self.assertRaises(ValueError, extractor, -1)
@@ -68,20 +58,24 @@ def _test_extractor(self, extractor, model):
 
  def _test_cam(self, name):
  # Get a pretrained model
- model = resnet18(pretrained=False).eval()
- conv_layer = 'layer4.1.relu'
+ model = mobilenet_v2(pretrained=False).eval()
+ conv_layer = None if name == "CAM" else 'features.16.conv.3'
 
+ kwargs = {}
+ # Speed up testing by reducing the number of samples
+ if name in ['SSCAM', 'ISCAM']:
+ kwargs['num_samples'] = 4
  # Hook the corresponding layer in the model
- extractor = cams.__dict__[name](model, conv_layer)
+ extractor = cams.__dict__[name](model, conv_layer, **kwargs)
 
  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.17.conv.3'
+ model = mobilenet_v2(pretrained=False).eval()
+ conv_layer = 'features.18.0'
 
  # Hook the corresponding layer in the model
  extractor = cams.__dict__[name](model, conv_layer)

torchcam/cams/gradcam.py

@@ -195,7 +195,7 @@ class SmoothGradCAMpp(_GradCAM):
  >>> from torchvision.models import resnet18
  >>> from torchcam.cams import SmoothGradCAMpp
  >>> model = resnet18(pretrained=True).eval()
- >>> cam = SmoothGradCAMpp(model, 'layer4', 'conv1')
+ >>> cam = SmoothGradCAMpp(model, 'layer4')
  >>> scores = model(input_tensor)
  >>> cam(class_idx=100)