Integrated Gradients

This notebook demonstrates how to use Integrated Gradients for attribution [1].

Setup

import importlib.util

IN_COLAB = importlib.util.find_spec("google.colab") is not None

if IN_COLAB:
    %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

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