torchxai

Explain any PyTorch vision model in one line of code.
10 methods. 24 models verified end-to-end. CNNs, ViTs, YOLO26, RF-DETR, DINOv2, and more.

Documentation • Quickstart • Installation • Methods • Examples • API Reference • Troubleshooting

Why torchxai?

Most explainability libraries make you write 20+ lines of boilerplate, only work with specific model architectures, and crash on Vision Transformers. torchxai does it in one line:

from torchxai import explain
heatmap = explain(model, "photo.jpg")

That's it. It works with any PyTorch vision model — ResNet, ViT, EfficientNet, CLIP, YOLO, Swin, DeiT, DINO, and more. No config files. No architecture-specific code. No boilerplate.

What makes it different

⚡ Quickstart

3 lines to your first explanation

from torchxai import explain
import torchvision.models as models

model = models.resnet50(pretrained=True)
heatmap = explain(model, "dog.jpg")  # Returns (224, 224) numpy array

Visualize it

from torchxai import explain, show_explanation
from PIL import Image

model = models.resnet50(pretrained=True)
heatmap = explain(model, "dog.jpg")
show_explanation(Image.open("dog.jpg"), heatmap, title="What the model sees")

Compare methods side by side

from torchxai import explain, create_comparison
from PIL import Image

heatmaps = {
    "GradCAM": explain(model, img, method="gradcam"),
    "EigenCAM": explain(model, img, method="eigencam"),
    "LayerCAM": explain(model, img, method="layercam"),
    "GradCAM++": explain(model, img, method="gradcam_pp"),
}
create_comparison(Image.open("dog.jpg"), heatmaps, save_path="comparison.png")

📦 Installation

pip install torchxai-explain

Requirements: Python 3.8+ and PyTorch 1.9+. That's all.

For development:

git clone https://github.com/rigvedrs/torchxai.git
cd torchxai
pip install -e ".[dev]"

🔬 Methods

torchxai includes 10 explainability methods. Each produces a saliency heatmap showing which image regions matter most for the model's prediction.

Don't know which to pick? Just use explain(model, image) — it auto-selects the best method for your model architecture.

→ Detailed method guide with math and intuition

🎯 Features

Accept any input format

# Tensor (already preprocessed)
heatmap = explain(model, tensor)

# PIL Image
heatmap = explain(model, Image.open("dog.jpg"))

# File path (string or Path)
heatmap = explain(model, "dog.jpg")
heatmap = explain(model, Path("dog.jpg"))

# 3D tensor (auto-adds batch dimension)
heatmap = explain(model, torch.randn(3, 224, 224))

Class-specific explanations

Ask "why did the model predict THIS class?" instead of just "where is it looking?"

heatmap_dog = explain(model, img, target_class=208)      # Labrador
heatmap_ball = explain(model, img, target_class=852)      # Tennis ball
heatmap_collar = explain(model, img, target_class=457)    # Bow tie

Customizable visualization

from torchxai import overlay_heatmap, show_explanation, save_heatmap

# Overlay on image
blended = overlay_heatmap(image, heatmap, colormap="jet", alpha=0.5)

# Side-by-side visualization
show_explanation(image, heatmap, save_path="output.png")

# Save raw heatmap
save_heatmap(heatmap, "heatmap.png", colormap="viridis")

Colormap options:

Transparency control:

Quantitative metrics

Don't just look at heatmaps — measure how good they are:

from torchxai import insertion_score, deletion_score, stability_score

# Does removing highlighted pixels drop confidence? (lower = better)
d_score = deletion_score(model, tensor, heatmap)

# Does revealing highlighted pixels restore confidence? (higher = better)
i_score = insertion_score(model, tensor, heatmap)

# Are explanations consistent under small perturbations? (higher = more stable)
s_score = stability_score(GradCAM(model), tensor)

Works on headless servers

No more TclError: no display name crashes. torchxai auto-detects headless environments (SSH, Docker, CI) and switches to the Agg matplotlib backend. Visualization functions save to files or return numpy arrays — no display required.

📖 Examples

CNN Explainability (ResNet, EfficientNet, VGG)

import torchvision.models as models
from torchxai import explain, show_explanation
from PIL import Image

# Any torchvision model
model = models.resnet50(pretrained=True)

# Explain
heatmap = explain(model, "cat.jpg")

# Visualize
show_explanation(Image.open("cat.jpg"), heatmap, save_path="resnet_explanation.png")

Vision Transformer (ViT, DeiT, Swin)

import timm
from torchxai import explain

# Any timm model — auto-detected
model = timm.create_model("vit_base_patch16_224", pretrained=True)
heatmap = explain(model, "photo.jpg")  # Just works

Custom target layer

from torchxai import GradCAM

# By name
cam = GradCAM(model, target_layer="layer3")

# By module reference
cam = GradCAM(model, target_layer=model.layer3[-1])

heatmap = cam("photo.jpg")

Batch comparison script

from torchxai import explain, create_comparison
from PIL import Image

model = ...  # Your model
image = Image.open("test.jpg")

results = {}
for method in ["gradcam", "eigencam", "layercam", "gradcam_pp"]:
    results[method] = explain(model, image, method=method)

create_comparison(image, results, save_path="all_methods.png")

→ More examples

📚 Documentation

🤝 Contributing

Contributions are welcome. See CONTRIBUTING.md for guidelines.

git clone https://github.com/rigvedrs/torchxai.git
cd torchxai
pip install -e ".[dev]"
pytest tests/ -v

📄 License

MIT License. See LICENSE for details.

⭐ Citation

If you use torchxai in your research, please cite:

@software{torchxai2026,
  title={torchxai: Universal Explainability for PyTorch Vision Models},
  author={Rigved Sandeep Shirvalkar},
  year={2026},
  url={https://github.com/rigvedrs/torchxai}
}

If torchxai helps your work, consider giving it a ⭐