Fix a type & refactor

.gitignore

@@ -100,4 +100,5 @@ ENV/
 # mypy
 .mypy_cache/
 
-.vscode
+.vscode
+.DS_Store

README.md

@@ -5,21 +5,23 @@ PyTorch implementation of [Grad-CAM (Gradient-weighted Class Activation Mapping)
 ## Dependencies
 * Python 2.7
 * PyTorch
+* torchvision
 
 ## Usage
 ```bash
-$ python main.py --image samples/cat_dog.jpg
+$ python main.py --image samples/cat_dog.jpg [--no-cuda]
 ```
 
 ## Examples
 ![](samples/cat_dog.png)
 
-||bull mastiff|tabby|
-|:-:|:-:|:-:|
-|**Grad-CAM [1]**|![](results/bull_mastiff_gcam.png)|![](results/tabby_gcam.png)|
-|Backpropogation|![](results/bull_mastiff_bp.png)|![](results/tabby_bp.png)|
-|Guided Backpropagation [2]|![](results/bull_mastiff_gbp.png)|![](results/tabby_gbp.png)|
-|**Guided Grad-CAM [1]**|![](results/bull_mastiff_ggcam.png)|![](results/tabby_ggcam.png)|
+||bull mastiff|tiger cat|boxer|
+|:-:|:-:|:-:|:-:|
+|Probability|0.54285|0.19302|0.10428|
+|**Grad-CAM [1]**|![](results/bull_mastiff_gcam.png)|![](results/tiger_cat_gcam.png)|![](results/boxer_gcam.png)|
+|Vanilla Backpropogation|![](results/bull_mastiff_bp.png)|![](results/tiger_cat_bp.png)|![](results/boxer_bp.png)|
+|Guided Backpropagation [2]|![](results/bull_mastiff_gbp.png)|![](results/tiger_cat_gbp.png)|![](results/boxer_gbp.png)|
+|**Guided Grad-CAM [1]**|![](results/bull_mastiff_ggcam.png)|![](results/tiger_cat_ggcam.png)|![](results/boxer_ggcam.png)|
 
 ## References
 \[1\] R. R. Selvaraju, A. Das, R. Vedantam, M. Cogswell, D. Parikh, and D. Batra. "Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization". arXiv, 2016<br>

gradcam.py → grad_cam.py

@@ -17,17 +17,15 @@
 from torch.nn import functional as F
 
 
-class PropagatationBase(object):
+class PropagationBase(object):
 
-    def __init__(self, model, target_layer, n_class, cuda):
+    def __init__(self, model, target_layer, cuda):
         self.model = model
         self.model.eval()
         self.cuda = cuda
         if self.cuda:
             self.model.cuda()
         self.target_layer = target_layer
-        self.n_class = n_class
-        self.probs = None
         self.all_fmaps = OrderedDict()
         self.all_grads = OrderedDict()
         self.set_hook_func()
@@ -36,28 +34,26 @@ def set_hook_func(self):
         raise NotImplementedError
 
     def encode_one_hot(self, idx):
-        one_hot = torch.FloatTensor(1, self.n_class).zero_()
+        one_hot = torch.FloatTensor(1, self.preds.size()[-1]).zero_()
         one_hot[0][idx] = 1.0
-        return one_hot
+        return one_hot.cuda() if self.cuda else one_hot
 
     def load_image(self, filename, transform):
         self.raw_image = cv2.imread(filename)[:, :, ::-1]
         self.raw_image = cv2.resize(self.raw_image, (224, 224))
-        self.image = transform(self.raw_image).unsqueeze(0)
-        if self.cuda:
-            self.image = self.image.cuda()
-        self.image = Variable(self.image, volatile=False, requires_grad=True)
+        image = transform(self.raw_image).unsqueeze(0)
+        image = image.cuda() if self.cuda else image
+        self.image = Variable(image, volatile=False, requires_grad=True)
 
     def forward(self):
         self.preds = self.model.forward(self.image)
         self.probs = F.softmax(self.preds)[0]
         self.prob, self.idx = self.probs.data.sort(0, True)
+        return self.prob, self.idx
 
     def backward(self, idx):
         self.model.zero_grad()
         one_hot = self.encode_one_hot(idx)
-        if self.cuda:
-            one_hot = one_hot.cuda()
         self.preds.backward(gradient=one_hot, retain_graph=True)
 
     def find(self, outputs, target_layer):
@@ -69,7 +65,7 @@ def find(self, outputs, target_layer):
         raise ValueError('invalid layer name: {}'.format(target_layer))
 
 
-class GradCAM(PropagatationBase):
+class GradCAM(PropagationBase):
 
     def set_hook_func(self):
 
@@ -117,7 +113,7 @@ def save(self, filename, gcam):
         cv2.imwrite(filename, gcam)
 
 
-class BackPropagation(PropagatationBase):
+class BackPropagation(PropagationBase):
 
     def set_hook_func(self):
 
@@ -147,7 +143,7 @@ def func_b(module, grad_in, grad_out):
 
             # Cut off negative gradients
             if isinstance(module, nn.ReLU):
-                return (F.threshold(grad_in[0], threshold=0.0, value=0.0),)
+                return (torch.clamp(grad_in[0], min=0.0),)
 
         for module in self.model.named_modules():
             module[1].register_backward_hook(func_b)

main.py

@@ -14,19 +14,18 @@
 import torchvision
 from torchvision import transforms
 
-from gradcam import BackPropagation, GradCAM, GuidedBackPropagation
+from grad_cam import BackPropagation, GradCAM, GuidedBackPropagation
 
 
 def main(args):
 
     # Load the synset words
-    file_name = 'synset_words.txt'
-    classes = list()
-    with open(file_name) as class_file:
-        for line in class_file:
-            classes.append(line.strip().split(' ', 1)[
-                           1].split(', ', 1)[0].replace(' ', '_'))
-
+    idx2cls = list()
+    with open('samples/synset_words.txt') as lines:
+        for line in lines:
+            line = line.strip().split(' ', 1)[1]
+            line = line.split(', ', 1)[0].replace(' ', '_')
+            idx2cls.append(line)
 
     print('Loading a model...')
     model = torchvision.models.resnet152(pretrained=True)
@@ -36,44 +35,44 @@ def main(args):
                              std=[0.229, 0.224, 0.225])
     ])
 
-
     print('\nGrad-CAM')
-    gcam = GradCAM(model=model, target_layer='layer4.2',
-                   n_class=1000, cuda=args.cuda)
+    gcam = GradCAM(model=model,
+                   target_layer='layer4.2',
+                   cuda=args.cuda)
     gcam.load_image(args.image, transform)
     gcam.forward()
 
-    for i in range(0, 5):
+    for i in range(0, 3):
         gcam.backward(idx=gcam.idx[i])
-        cls_name = classes[gcam.idx[i]]
+        cls_name = idx2cls[gcam.idx[i]]
         output = gcam.generate()
         print('\t{:.5f}\t{}'.format(gcam.prob[i], cls_name))
         gcam.save('results/{}_gcam.png'.format(cls_name), output)
 
-
-    print('\nBackpropagation')
-    bp = BackPropagation(model=model, target_layer='conv1',
-                         n_class=1000, cuda=args.cuda)
+    print('\nVanilla Backpropagation')
+    bp = BackPropagation(model=model,
+                         target_layer='conv1',
+                         cuda=args.cuda)
     bp.load_image(args.image, transform)
     bp.forward()
 
-    for i in range(0, 5):
+    for i in range(0, 3):
         bp.backward(idx=bp.idx[i])
-        cls_name = classes[bp.idx[i]]
+        cls_name = idx2cls[bp.idx[i]]
         output = bp.generate()
         print('\t{:.5f}\t{}'.format(bp.prob[i], cls_name))
         bp.save('results/{}_bp.png'.format(cls_name), output)
 
-
     print('\nGuided Backpropagation')
-    gbp = GuidedBackPropagation(model=model, target_layer='conv1',
-                                n_class=1000, cuda=args.cuda)
+    gbp = GuidedBackPropagation(model=model,
+                                target_layer='conv1',
+                                cuda=args.cuda)
     gbp.load_image(args.image, transform)
     gbp.forward()
 
-    for i in range(0, 5):
+    for i in range(0, 3):
         cls_idx = gcam.idx[i]
-        cls_name = classes[cls_idx]
+        cls_name = idx2cls[cls_idx]
 
         gcam.backward(idx=cls_idx)
         output_gcam = gcam.generate()