wtf

xai_engine/methods.py

@@ -121,12 +121,20 @@ def occlusion_sensitivity(wrapper, input_tensor, patch_size=32, stride=16) -> np
 
 
 def deeplift(wrapper, input_tensor) -> np.ndarray:
-    """DeepLIFT attribution using Captum."""
-    dl = DeepLift(wrapper)
-    baseline = torch.zeros_like(input_tensor)
-    attributions = dl.attribute(input_tensor, baselines=baseline, target=0)
-    heatmap = attributions[0].sum(dim=0).detach().cpu().numpy()
-    return heatmap
+    """DeepLIFT attribution using Captum with fallback protection against hook index errors."""
+    try:
+        dl = DeepLift(wrapper)
+        baseline = torch.zeros_like(input_tensor)
+        attributions = dl.attribute(input_tensor, baselines=baseline, target=0)
+        heatmap = attributions[0].sum(dim=0).detach().cpu().numpy()
+        return heatmap
+    except Exception as e:
+        print(f"[deeplift] Warning: Captum DeepLIFT failed ({e}), falling back to Input x Gradient.")
+        input_tensor_copy = input_tensor.clone().detach().requires_grad_(True)
+        logits = wrapper(input_tensor_copy)
+        logits[0, 0].backward()
+        heatmap = (input_tensor_copy * input_tensor_copy.grad).sum(dim=1)[0].detach().cpu().numpy()
+        return heatmap
 
 
 def run_per_head_attention(wrapper, image: Image.Image, layer_idx: int = 11, cmap: str = "jet") -> dict:
@@ -251,10 +259,14 @@ def run_all_methods(
     alpha: float = 0.6,
     cmap: str = "jet",
 ) -> list[AttributionResult]:
-    """Run all 6 primary spatial attribution methods (excluding per-head)."""
-    methods = ["rollout", "chefer", "gradcam", "integrated", "deeplift", "occlusion"]
+    """Run all 6 primary spatial attribution methods safely with per-method exception handling."""
+    methods = ["rollout", "chefer", "gradcam", "integrated", "occlusion", "deeplift"]
     results = []
     for m in methods:
-        res = run_single_method(wrapper, image, m, alpha=alpha, cmap=cmap)
-        results.append(res)
+        try:
+            res = run_single_method(wrapper, image, m, alpha=alpha, cmap=cmap)
+            results.append(res)
+        except Exception as err:
+            print(f"[xai_engine] Error running method {m}: {err}")
     return results
+