Skip to main content

Explainability and saliency maps for PyTorch vision models — GradCAM, EigenCAM, Attention Rollout for CNNs, Vision Transformers (ViT), CLIP, YOLO, DETR, DINO, and SAM. One-line API with auto-detection.

Project description

torchxai — One line to explain any vision model

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.

PyPI Python License Stars

DocumentationQuickstartInstallationMethodsExamplesAPI ReferenceTroubleshooting


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.

Quick-start workflow

What makes it different

Feature torchxai pytorch-grad-cam captum tf-explain
One-line API explain(model, img) ❌ 15+ lines ❌ 20+ lines ❌ TF only
Auto-detects architecture ✅ CNN + ViT + YOLO ❌ Manual config ❌ Manual config
Vision Transformer support ✅ Native ⚠️ Partial ⚠️ Partial
String path input explain(m, "dog.jpg")
Headless server support ✅ Auto Agg backend ❌ Crashes
Built-in metrics ✅ Insertion/Deletion/Stability ⚠️ Separate
Graceful fallbacks ✅ Always produces output ❌ Crashes ❌ Crashes

⚡ 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")

Explanation output

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")

Method comparison


📦 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.

Method Type Needs Gradients Best For
GradCAM Gradient-weighted General-purpose CNN explanation
GradCAM++ Gradient-weighted Multiple objects of same class
LayerCAM Gradient-weighted Fine-grained spatial detail
EigenCAM Activation-based Fast, gradient-free explanation
ScoreCAM Perturbation-based Most faithful, no gradients
SmoothGrad Input gradients Pixel-level attribution
Integrated Gradients Axiomatic Theoretically grounded
RISE Black-box sampling Model-agnostic, any architecture
Attention Rollout Attention-based Vision Transformers (ViT)
Transformer Attribution Gradient + Attention Class-specific ViT explanation

All methods compared

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

Class-specific explanations

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:

Colormap options

Transparency control:

Transparency options

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

Document Description
Getting Started Installation, first explanation in 60 seconds
API Reference Every function, every parameter, return types
Methods Guide How each method works, when to use which
Advanced Usage Custom layers, metrics, deployment, optimization
Troubleshooting Common errors with exact fixes
Model Compatibility Tested models with proof images
Contributing How to contribute

🤝 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 ⭐

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

torchxai_explain-0.3.0.tar.gz (65.4 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

torchxai_explain-0.3.0-py3-none-any.whl (65.7 kB view details)

Uploaded Python 3

File details

Details for the file torchxai_explain-0.3.0.tar.gz.

File metadata

  • Download URL: torchxai_explain-0.3.0.tar.gz
  • Upload date:
  • Size: 65.4 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for torchxai_explain-0.3.0.tar.gz
Algorithm Hash digest
SHA256 571c7ac3e16a6a113d44f5b897984c035e5f2274a1039df6ca4b1e3c506a7f3d
MD5 b98a4049a0ddb8488a71e1f91a92f64f
BLAKE2b-256 28fc61a18f8afdad66314674c89a58a8edb80f65555a65ebac14916bfd740845

See more details on using hashes here.

Provenance

The following attestation bundles were made for torchxai_explain-0.3.0.tar.gz:

Publisher: publish.yml on rigvedrs/torchxai

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file torchxai_explain-0.3.0-py3-none-any.whl.

File metadata

File hashes

Hashes for torchxai_explain-0.3.0-py3-none-any.whl
Algorithm Hash digest
SHA256 0a24b7bd34aec90eef9c568887aa2a7db5faae9067020465ea3fd676cfd9fa89
MD5 4e31ac97eb20c9796e48168e27d5b68a
BLAKE2b-256 453f709f8e0d91c4de1431cacd0010b89ade75aa1c0dbd4c4b48477060ba18b9

See more details on using hashes here.

Provenance

The following attestation bundles were made for torchxai_explain-0.3.0-py3-none-any.whl:

Publisher: publish.yml on rigvedrs/torchxai

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page