Merge branch 'master' of github.com:jacobgil/pytorch-grad-cam into master

Author: jacobgil

README.md

@@ -282,7 +282,7 @@ two smoothing methods are supported:
 Usage: `python cam.py --image-path <path_to_image> --method <method> --output-dir <output_dir_path> `
 
 
-To use with a specific device, like cpu, cuda, cuda:0 or mps:
+To use with a specific device, like cpu, cuda, cuda:0, mps or hpu:
 `python cam.py --image-path <path_to_image> --device cuda  --output-dir <output_dir_path> `
 
 ----------

cam.py

@@ -78,6 +78,9 @@ def get_args():
         'kpcacam': KPCA_CAM
     }
 
+    if args.device=='hpu':
+        import habana_frameworks.torch.core as htcore
+
     model = models.resnet50(pretrained=True).to(torch.device(args.device)).eval()
 
     # Choose the target layer you want to compute the visualization for.

pytorch_grad_cam/base_cam.py

@@ -25,6 +25,13 @@ def __init__(
 
         # Use the same device as the model.
         self.device = next(self.model.parameters()).device
+        if 'hpu' in str(self.device):
+            try:
+                import habana_frameworks.torch.core as htcore
+            except ImportError as error:
+                error.msg = f"Could not import habana_frameworks.torch.core. {error.msg}."
+                raise error
+            self.__htcore = htcore
         self.reshape_transform = reshape_transform
         self.compute_input_gradient = compute_input_gradient
         self.uses_gradients = uses_gradients
@@ -97,6 +104,8 @@ def forward(
             self.model.zero_grad()
             loss = sum([target(output) for target, output in zip(targets, outputs)])
             loss.backward(retain_graph=True)
+            if 'hpu' in str(self.device):
+                self.__htcore.mark_step()
 
         # In most of the saliency attribution papers, the saliency is
         # computed with a single target layer.

pytorch_grad_cam/utils/model_targets.py

@@ -59,13 +59,9 @@ class SemanticSegmentationTarget:
     def __init__(self, category, mask):
         self.category = category
         self.mask = torch.from_numpy(mask)
-        if torch.cuda.is_available():
-            self.mask = self.mask.cuda()
-        if torch.backends.mps.is_available():
-            self.mask = self.mask.to("mps")
 
     def __call__(self, model_output):
-        return (model_output[self.category, :, :] * self.mask).sum()
+        return (model_output[self.category, :, :] * self.mask.to(model_output.device)).sum()
 
 
 class FasterRCNNBoxScoreTarget: