From 0f3123bd7565aa4d77cdbeb8dde1c8f7e6e1b734 Mon Sep 17 00:00:00 2001 From: Daniel Deleon Date: Mon, 16 Dec 2024 11:36:11 -0800 Subject: [PATCH 1/2] add gaudi hpu option Signed-off-by: Daniel Deleon --- README.md | 2 +- cam.py | 3 +++ pytorch_grad_cam/base_cam.py | 5 +++++ 3 files changed, 9 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 95e48c272..2b588dde4 100644 --- a/README.md +++ b/README.md @@ -281,7 +281,7 @@ two smoothing methods are supported: Usage: `python cam.py --image-path --method --output-dir ` -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 --device cuda --output-dir ` ---------- diff --git a/cam.py b/cam.py index 9f3e377c1..935afeb22 100644 --- a/cam.py +++ b/cam.py @@ -77,6 +77,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. diff --git a/pytorch_grad_cam/base_cam.py b/pytorch_grad_cam/base_cam.py index 4b2850a98..78225822d 100644 --- a/pytorch_grad_cam/base_cam.py +++ b/pytorch_grad_cam/base_cam.py @@ -25,6 +25,9 @@ def __init__( # Use the same device as the model. self.device = next(self.model.parameters()).device + if 'hpu' in str(self.device): + import habana_frameworks.torch.core as htcore + self.__htcore = htcore self.reshape_transform = reshape_transform self.compute_input_gradient = compute_input_gradient self.uses_gradients = uses_gradients @@ -97,6 +100,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. From 6a8e588685a8400c1da53e43064861b012a4ee6c Mon Sep 17 00:00:00 2001 From: Daniel Deleon Date: Mon, 16 Dec 2024 14:21:19 -0800 Subject: [PATCH 2/2] add try except block Signed-off-by: Daniel Deleon --- pytorch_grad_cam/base_cam.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pytorch_grad_cam/base_cam.py b/pytorch_grad_cam/base_cam.py index 78225822d..44ae5b906 100644 --- a/pytorch_grad_cam/base_cam.py +++ b/pytorch_grad_cam/base_cam.py @@ -26,7 +26,11 @@ def __init__( # Use the same device as the model. self.device = next(self.model.parameters()).device if 'hpu' in str(self.device): - import habana_frameworks.torch.core as htcore + 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