Integrated Gradients

This notebook demonstrates how to use Integrated Gradients for attribution.

Setup

import importlib.util

DEV = True

if importlib.util.find_spec("google.colab") is not None:
    MODE = "colab-dev" if DEV else "colab"
else:
    MODE = "local"

if MODE == "colab":
    %pip install -q tdhook
elif MODE == "colab-dev":
    !rm -r tdhook
    !git clone https://github.com/Xmaster6y/tdhook -b main
    %pip install -q ./tdhook

Usage

Load model and prepare image

import torch
import timm
from PIL import Image
from tensordict import TensorDict
from tdhook.attribution import IntegratedGradients

/home/docs/checkouts/readthedocs.org/user_builds/tdhook/checkouts/v0.1.2/.venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

model = timm.create_model("vgg16.tv_in1k", pretrained=True)
data_config = timm.data.resolve_model_data_config(model)
transforms = timm.data.create_transform(**data_config, is_training=False)

image = Image.open("../assets/zebra_1.jpg").convert("RGB")
image_tensor = transforms(image)

Define attribution target (zebra class = 340)

def init_attr_targets(targets, _):
    zebra_logit = targets["output"][..., 340]
    return TensorDict(out=zebra_logit, batch_size=targets.batch_size)

Compute attribution

with IntegratedGradients(init_attr_targets=init_attr_targets).prepare(model) as hooked_model:
    td = TensorDict({"input": image_tensor, ("baseline", "input"): torch.zeros_like(image_tensor)}).unsqueeze(0)
    td = hooked_model(td)

Access attribution:

td.get(("attr", "input")).shape

torch.Size([1, 3, 224, 224])