Adding capability to choose device other than cpu and fixing/generalize

src/gradcam.py

@@ -18,6 +18,17 @@ def __init__(self, model, target_layer):
         self.model = model
         self.target_layer = target_layer
         self.gradients = None
+        self.conv_output = None
+        self.forward_hook()
+
+    def hook_Func(self, module, input, output):
+        self.conv_output = output
+        self.conv_output.register_hook(self.save_gradient)
+
+    def forward_hook(self):
+        for module_pos, module in self.model._modules.items():
+            if module_pos == self.target_layer:
+                module.register_forward_hook(self.hook_Func)
 
     def save_gradient(self, grad):
         self.gradients = grad
@@ -26,13 +37,9 @@ def forward_pass_on_convolutions(self, x):
         """
             Does a forward pass on convolutions, hooks the function at given layer
         """
-        conv_output = None
-        for module_pos, module in self.model.features._modules.items():
-            x = module(x)  # Forward
-            if int(module_pos) == self.target_layer:
-                x.register_hook(self.save_gradient)
-                conv_output = x  # Save the convolution output on that layer
-        return conv_output, x
+        output = None
+        output = self.model(x)
+        return self.conv_output, output
 
     def forward_pass(self, x):
         """
@@ -50,8 +57,9 @@ class GradCam():
     """
         Produces class activation map
     """
-    def __init__(self, model, target_layer):
-        self.model = model
+    def __init__(self, model, target_layer,device_id='cpu'):
+        self.device = device_id
+        self.model = model.to(self.device)
         self.model.eval()
         # Define extractor
         self.extractor = CamExtractor(self.model, target_layer)
@@ -60,21 +68,21 @@ def generate_cam(self, input_image, target_class=None):
         # Full forward pass
         # conv_output is the output of convolutions at specified layer
         # model_output is the final output of the model (1, 1000)
-        conv_output, model_output = self.extractor.forward_pass(input_image)
+        conv_output, model_output = self.extractor.forward_pass(input_image.to(self.device))
         if target_class is None:
-            target_class = np.argmax(model_output.data.numpy())
+            target_class = np.argmax(model_output.data.cpu().numpy())
         # Target for backprop
         one_hot_output = torch.FloatTensor(1, model_output.size()[-1]).zero_()
         one_hot_output[0][target_class] = 1
         # Zero grads
         self.model.features.zero_grad()
         self.model.classifier.zero_grad()
         # Backward pass with specified target
-        model_output.backward(gradient=one_hot_output, retain_graph=True)
+        model_output.backward(gradient=one_hot_output.to(self.device), retain_graph=True)
         # Get hooked gradients
-        guided_gradients = self.extractor.gradients.data.numpy()[0]
+        guided_gradients = self.extractor.gradients.data.cpu().numpy()[0]
         # Get convolution outputs
-        target = conv_output.data.numpy()[0]
+        target = conv_output.data.cpu().numpy()[0]
         # Get weights from gradients
         weights = np.mean(guided_gradients, axis=(1, 2))  # Take averages for each gradient
         # Create empty numpy array for cam

src/guided_backprop.py

@@ -16,8 +16,9 @@ class GuidedBackprop():
     """
        Produces gradients generated with guided back propagation from the given image
     """
-    def __init__(self, model):
-        self.model = model
+    def __init__(self, model, device_id = "cpu"):
+        self.device = device_id
+        self.model = model.to(self.device)
         self.gradients = None
         self.forward_relu_outputs = []
         # Put model in evaluation mode
@@ -63,17 +64,17 @@ def relu_forward_hook_function(module, ten_in, ten_out):
 
     def generate_gradients(self, input_image, target_class):
         # Forward pass
-        model_output = self.model(input_image)
+        model_output = self.model(input_image.to(self.device))
         # Zero gradients
         self.model.zero_grad()
         # Target for backprop
         one_hot_output = torch.FloatTensor(1, model_output.size()[-1]).zero_()
         one_hot_output[0][target_class] = 1
         # Backward pass
-        model_output.backward(gradient=one_hot_output)
+        model_output.backward(gradient=one_hot_output.to(self.device))
         # Convert Pytorch variable to numpy array
         # [0] to get rid of the first channel (1,3,224,224)
-        gradients_as_arr = self.gradients.data.numpy()[0]
+        gradients_as_arr = self.gradients.data.cpu().numpy()[0]
         return gradients_as_arr
 
 

src/integrated_gradients.py

@@ -13,8 +13,9 @@ class IntegratedGradients():
     """
         Produces gradients generated with integrated gradients from the image
     """
-    def __init__(self, model):
-        self.model = model
+    def __init__(self, model, device_id='cpu'):
+        self.device = device_id
+        self.model = model.to(self.device)
         self.gradients = None
         # Put model in evaluation mode
         self.model.eval()
@@ -38,17 +39,17 @@ def generate_images_on_linear_path(self, input_image, steps):
 
     def generate_gradients(self, input_image, target_class):
         # Forward
-        model_output = self.model(input_image)
+        model_output = self.model(input_image.to(self.device))
         # Zero grads
         self.model.zero_grad()
         # Target for backprop
         one_hot_output = torch.FloatTensor(1, model_output.size()[-1]).zero_()
         one_hot_output[0][target_class] = 1
         # Backward pass
-        model_output.backward(gradient=one_hot_output)
+        model_output.backward(gradient=one_hot_output.to(self.device))
         # Convert Pytorch variable to numpy array
         # [0] to get rid of the first channel (1,3,224,224)
-        gradients_as_arr = self.gradients.data.numpy()[0]
+        gradients_as_arr = self.gradients.data.cpu().numpy()[0]
         return gradients_as_arr
 
     def generate_integrated_gradients(self, input_image, target_class, steps):

src/layer_activation_with_guided_backprop.py

@@ -16,21 +16,27 @@ class GuidedBackprop():
     """
        Produces gradients generated with guided back propagation from the given image
     """
-    def __init__(self, model):
-        self.model = model
+    def __init__(self, model, cnn_layer, device_id= 'cpu'):
+        self.device = device_id
+        self.model = model.to(device)
         self.gradients = None
         self.forward_relu_outputs = []
         # Put model in evaluation mode
+        self.layer = cnn_layer
+        self.conv_output = None
         self.model.eval()
         self.update_relus()
         self.hook_layers()
 
     def hook_layers(self):
         def hook_function(module, grad_in, grad_out):
             self.gradients = grad_in[0]
+        def hook_function_forward(module, input, output):
+            self.conv_output = output
         # Register hook to the first layer
         first_layer = list(self.model.features._modules.items())[0][1]
         first_layer.register_backward_hook(hook_function)
+        self.layer.register_forward_hook(hook_function_forward)
 
     def update_relus(self):
         """
@@ -64,22 +70,14 @@ def relu_forward_hook_function(module, ten_in, ten_out):
     def generate_gradients(self, input_image, target_class, cnn_layer, filter_pos):
         self.model.zero_grad()
         # Forward pass
-        x = input_image
-        for index, layer in enumerate(self.model.features):
-            # Forward pass layer by layer
-            # x is not used after this point because it is only needed to trigger
-            # the forward hook function
-            x = layer(x)
-            # Only need to forward until the selected layer is reached
-            if index == cnn_layer:
-                # (forward hook function triggered)
-                break
-        conv_output = torch.sum(torch.abs(x[0, filter_pos]))
+        x = input_image.to(self.device)
+        op = self.model(x)
+        conv_output = torch.sum(torch.abs(self.conv_output[0, filter_pos]))
         # Backward pass
         conv_output.backward()
         # Convert Pytorch variable to numpy array
         # [0] to get rid of the first channel (1,3,224,224)
-        gradients_as_arr = self.gradients.data.numpy()[0]
+        gradients_as_arr = self.gradients.data.cpu().numpy()[0]
         return gradients_as_arr
 
 
@@ -93,7 +91,7 @@ def generate_gradients(self, input_image, target_class, cnn_layer, filter_pos):
     # File export name
     file_name_to_export = file_name_to_export + '_layer' + str(cnn_layer) + '_filter' + str(filter_pos)
     # Guided backprop
-    GBP = GuidedBackprop(pretrained_model)
+    GBP = GuidedBackprop(pretrained_model, cnn_layer)
     # Get gradients
     guided_grads = GBP.generate_gradients(prep_img, target_class, cnn_layer, filter_pos)
     # Save colored gradients

src/scorecam.py

@@ -18,18 +18,24 @@ class CamExtractor():
     def __init__(self, model, target_layer):
         self.model = model
         self.target_layer = target_layer
+        self.layer_output = None
+        self.forward_hook()
+
+    def hook_Func(self, module, input, output):
+        self.layer_output = output
+
+    def forward_hook(self):
+        for module_pos, module in self.model._modules.items():
+            if module_pos == self.target_layer:
+                module.register_forward_hook(self.hook_Func)
 
     def forward_pass_on_convolutions(self, x):
         """
             Does a forward pass on convolutions, hooks the function at given layer
         """
-        conv_output = None
-        for module_pos, module in self.model.features._modules.items():
-            x = module(x)  # Forward
-            if int(module_pos) == self.target_layer:
-                conv_output = x  # Save the convolution output on that layer
-        return conv_output, x
-
+        output = None
+        output = = self.model(x)
+        return self.layer_output, output
     def forward_pass(self, x):
         """
             Does a full forward pass on the model
@@ -46,8 +52,9 @@ class ScoreCam():
     """
         Produces class activation map
     """
-    def __init__(self, model, target_layer):
-        self.model = model
+    def __init__(self, model, target_layer, device_id= 'cpu'):
+        self.device = device_id
+        self.model = model.to(self.device)
         self.model.eval()
         # Define extractor
         self.extractor = CamExtractor(self.model, target_layer)
@@ -56,9 +63,9 @@ def generate_cam(self, input_image, target_class=None):
         # Full forward pass
         # conv_output is the output of convolutions at specified layer
         # model_output is the final output of the model (1, 1000)
-        conv_output, model_output = self.extractor.forward_pass(input_image)
+        conv_output, model_output = self.extractor.forward_pass(input_image.to(self.device))
         if target_class is None:
-            target_class = np.argmax(model_output.data.numpy())
+            target_class = np.argmax(model_output.data.cpu().numpy())
         # Get convolution outputs
         target = conv_output[0]
         # Create empty numpy array for cam
@@ -74,8 +81,8 @@ def generate_cam(self, input_image, target_class=None):
             # Scale between 0-1
             norm_saliency_map = (saliency_map - saliency_map.min()) / (saliency_map.max() - saliency_map.min())
             # Get the target score
-            w = F.softmax(self.extractor.forward_pass(input_image*norm_saliency_map)[1],dim=1)[0][target_class]
-            cam += w.data.numpy() * target[i, :, :].data.numpy()
+            w = F.softmax(self.extractor.forward_pass(input_image.to(self.device)*norm_saliency_map)[1],dim=1)[0][target_class]
+            cam += w.data.cpu().numpy() * target[i, :, :].data.cpu().numpy()
         cam = np.maximum(cam, 0)
         cam = (cam - np.min(cam)) / (np.max(cam) - np.min(cam))  # Normalize between 0-1
         cam = np.uint8(cam * 255)  # Scale between 0-255 to visualize

src/vanilla_backprop.py

@@ -12,8 +12,9 @@ class VanillaBackprop():
     """
         Produces gradients generated with vanilla back propagation from the image
     """
-    def __init__(self, model):
-        self.model = model
+    def __init__(self, model, device_id ='cpu'):
+        self.device = device_id
+        self.model = model.to(self.device)
         self.gradients = None
         # Put model in evaluation mode
         self.model.eval()
@@ -30,17 +31,17 @@ def hook_function(module, grad_in, grad_out):
 
     def generate_gradients(self, input_image, target_class):
         # Forward
-        model_output = self.model(input_image)
+        model_output = self.model(input_image.to(self.device))
         # Zero grads
         self.model.zero_grad()
         # Target for backprop
         one_hot_output = torch.FloatTensor(1, model_output.size()[-1]).zero_()
         one_hot_output[0][target_class] = 1
         # Backward pass
-        model_output.backward(gradient=one_hot_output)
+        model_output.backward(gradient=one_hot_output.to(self.device))
         # Convert Pytorch variable to numpy array
         # [0] to get rid of the first channel (1,3,224,224)
-        gradients_as_arr = self.gradients.data.numpy()[0]
+        gradients_as_arr = self.gradients.data.cpu().numpy()[0]
         return gradients_as_arr