Spaces:
Sleeping
Sleeping
File size: 5,128 Bytes
8c58a75 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """
SmoothGrad: Removing Noise from Gradients
Paper: https://arxiv.org/abs/1706.03825 (Smilkov et al., 2017)
Vanilla gradients of the output w.r.t. the input are noisy because the gradient
function is highly non-linear in practice. SmoothGrad averages gradients computed
on noisy copies of the input to denoise the attribution map.
Formula:
SmoothGrad(x) = (1/N) x Σ_k ∂F(x + ε_k) / ∂x
where ε_k ~ N(0, σ²)
σ is typically set as a fraction of the input range (0.1-0.2 x (max - min)).
Variants included:
- SmoothGrad (vanilla)
- SmoothGrad-Squared (SG-SQ): emphasizes strongest signals
- SmoothGrad-VAR (SG-VAR): highlights where gradients are consistent
"""
import numpy as np
import torch
import torch.nn.functional as F
from typing import Optional, Tuple
class SmoothGrad:
"""
SmoothGrad noise-averaged gradient attribution.
Args:
model: PyTorch model
n_samples: Number of noisy samples to average (50 typical)
noise_level: Noise std as fraction of input value range (default 0.15)
variant: 'standard' | 'squared' | 'var'
"""
def __init__(
self,
model: torch.nn.Module,
n_samples: int = 50,
noise_level: float = 0.15,
variant: str = "standard",
):
self.model = model
self.n_samples = n_samples
self.noise_level = noise_level
self.variant = variant
self.model.eval()
def __call__(
self,
input_tensor: torch.Tensor,
class_idx: Optional[int] = None,
) -> Tuple[np.ndarray, np.ndarray]:
"""
Compute SmoothGrad attribution map.
Args:
input_tensor: (1, C, H, W) preprocessed image tensor
class_idx: Target class. If None, uses the argmax prediction.
Returns:
smooth_attrs: (C, H, W) numpy array — channel-wise attributions
smooth_map: (H, W) numpy array in [0, 1] — collapsed saliency map
"""
with torch.no_grad():
logits = self.model(input_tensor)
if class_idx is None:
class_idx = logits.argmax(dim=1).item()
# Noise standard deviation proportional to input value range
val_range = input_tensor.max().item() - input_tensor.min().item()
sigma = self.noise_level * val_range
device = input_tensor.device
all_grads = []
for _ in range(self.n_samples):
noise = torch.randn_like(input_tensor) * sigma
noisy_input = (input_tensor.detach() + noise).requires_grad_(True)
logits = self.model(noisy_input)
score = logits[0, class_idx]
self.model.zero_grad()
score.backward(retain_graph=False)
grad = noisy_input.grad.detach().cpu()
all_grads.append(grad.squeeze(0))
# Stack: (n_samples, C, H, W)
grads_stack = torch.stack(all_grads, dim=0)
if self.variant == "standard":
# Average of gradients
smooth_attrs = grads_stack.mean(dim=0)
elif self.variant == "squared":
# Average of squared gradients — amplifies confident attributions
smooth_attrs = (grads_stack ** 2).mean(dim=0)
elif self.variant == "var":
# Variance of gradients — highlights regions where model is certain
smooth_attrs = grads_stack.var(dim=0)
else:
smooth_attrs = grads_stack.mean(dim=0)
smooth_attrs_np = smooth_attrs.numpy()
# Collapse to single map: take absolute value, then mean across channels
smooth_map = np.abs(smooth_attrs_np).mean(axis=0)
smooth_map = self._normalize(smooth_map)
return smooth_attrs_np, smooth_map
@staticmethod
def _normalize(arr: np.ndarray) -> np.ndarray:
min_val, max_val = arr.min(), arr.max()
if max_val - min_val < 1e-8:
return np.zeros_like(arr)
return (arr - min_val) / (max_val - min_val)
class VanillaGradients:
"""
Vanilla gradient saliency map (baseline comparison).
∂F(x) / ∂x — how much does each pixel affect the prediction?
"""
def __init__(self, model: torch.nn.Module):
self.model = model
self.model.eval()
def __call__(
self,
input_tensor: torch.Tensor,
class_idx: Optional[int] = None,
) -> Tuple[np.ndarray, np.ndarray]:
input_tensor = input_tensor.requires_grad_(True)
logits = self.model(input_tensor)
if class_idx is None:
class_idx = logits.argmax(dim=1).item()
self.model.zero_grad()
logits[0, class_idx].backward()
grads = input_tensor.grad.detach().cpu().squeeze(0).numpy()
saliency = np.abs(grads).mean(axis=0)
saliency = self._normalize(saliency)
return grads, saliency
@staticmethod
def _normalize(arr: np.ndarray) -> np.ndarray:
min_val, max_val = arr.min(), arr.max()
if max_val - min_val < 1e-8:
return np.zeros_like(arr)
return (arr - min_val) / (max_val - min_val)
|