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

Author: frgfm

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)