meta-pytorch/captum

What is Captum?

Captum is an open‑source library that adds model‑interpretability tools to PyTorch. It implements a collection of attribution algorithms (Integrated Gradients, DeepLift, Gradient Shap, SmoothGrad, TCAV, TracIn, etc.) that let you ask “which input features, neurons, or training examples contributed to this prediction?” The library works directly with any PyTorch model, including those built with torchvision, torchtext, and other domain‑specific extensions.

Who is it for?

Model developers who want to debug or improve their networks by seeing what drives predictions.
Researchers building new interpretability methods who need a benchmark suite.
Production engineers who need to generate explanations for end‑users (e.g., why a recommendation was made).

Quick installation

# from PyPI (most common)
pip install captum

# or via conda
conda install -c pytorch captum   # or conda‑forge channel

For the latest development version:

git clone https://github.com/pytorch/captum.git
cd captum
pip install -e .

(Use -e .[dev] or -e .[tutorials] to add extra dev or tutorial dependencies.)

Minimal example (the README’s toy model)

import torch, torch.nn as nn
from captum.attr import IntegratedGradients, GradientShap, DeepLift, NoiseTunnel

class ToyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin1 = nn.Linear(3, 3)
        self.relu = nn.ReLU()
        self.lin2 = nn.Linear(3, 2)
        # deterministic weights for the demo
        self.lin1.weight = nn.Parameter(torch.arange(-4., 5.).view(3,3))
        self.lin1.bias   = nn.Parameter(torch.zeros(1,3))
        self.lin2.weight = nn.Parameter(torch.arange(-3., 3.).view(2,3))
        self.lin2.bias   = nn.Parameter(torch.ones(1,2))
    def forward(self, x):
        return self.lin2(self.relu(self.lin1(x)))

model = ToyModel().eval()
input   = torch.rand(2, 3)
baseline = torch.zeros(2, 3)

# Integrated Gradients
ig = IntegratedGradients(model)
attr, delta = ig.attribute(input, baseline, target=0, return_convergence_delta=True)
print('IG attributions:', attr)
print('Delta:', delta)

# GradientShap (uses a baseline distribution)
gs = GradientShap(model)
baseline_dist = torch.randn(10, 3) * 0.001
attr, delta = gs.attribute(input, stdevs=0.09, n_samples=4,
                           baselines=baseline_dist, target=0,
                           return_convergence_delta=True)
print('GradShap attributions:', attr)

# Smoothing with NoiseTunnel (SmoothGrad)
nt = NoiseTunnel(IntegratedGradients(model))
attr, delta = nt.attribute(input, nt_type='smoothgrad', stdevs=0.02,
                           nt_samples=4, baselines=baseline, target=0,
                           return_convergence_delta=True)
print('SmoothGrad IG:', attr)

The code prints per‑feature attribution scores (positive → supports the prediction, negative → opposes) and a convergence delta that measures how accurate the integral approximation is.

Main capabilities (as listed in the README)

Attribution methods: Integrated Gradients, DeepLift, Gradient Shap, SmoothGrad/VarGrad, TCAV, TracIn, etc.
Neuron‑ and layer‑level analysis: NeuronConductance, LayerConductance let you see which internal units matter.
Counterfactual & adversarial utilities: minimal input perturbations for explanations or robustness testing.
Compatibility: works with any PyTorch model, including torchvision, torchtext, and custom architectures.
Bench‑marking: researchers can compare new algorithms against the built‑in suite.

Where to learn more

Official docs: https://captum.ai/
Tutorials (install with pip install -e .[tutorials])
FAQ: docs/faq.md
Talks & papers linked in the README (NeurIPS 2019, KDD 2020, etc.)

Bottom line: Captum is the go‑to library for anyone who needs to explain PyTorch models, whether for research, debugging, or production‑grade user‑facing explanations.