mkurman Overview
Captum (Comprehension in PyTorch) provides model interpretability for PyTorch models. Implements Integrated Gradients, Gradient SHAP, DeepLIFT, Occlusion, Feature Ablation, and Layer Conductance. Supports computer vision, NLP, and tabular models.
Installation
uv pip install captum
Integrated Gradients
import torch
import torch.nn as nn
from captum.attr import IntegratedGradients
model = nn.Linear(10, 2)
input = torch.randn(1, 10)
baseline = torch.zeros(1, 10)
ig = IntegratedGradients(model)
attrs = ig.attribute(input, baseline, target=0)
print(f"Feature attributions: {attrs}")
Occlusion
from captum.attr import Occlusion
occ = Occlusion(model)
attrs = occ.attribute(input, target=0, sliding_window_shapes=(1,)) # 1D
print(attrs)
Visualization
from captum.attr import visualization as viz
_ = viz.visualize_image_attr(
attrs.squeeze().numpy(),
original_image=input.squeeze().numpy(),
method="heat_map",
sign="absolute_value",
show_colorbar=True,
)