perception_models/LEGRAD_PE_USAGE.md

LeGrad + PE Perception Encoder Notebook Usage

This repository includes a notebook legrad_perception_encoder.ipynb that demonstrates how to run LeGrad explanations on the PE CoCa-style vision encoder.

1. Environment and installation

• Install this repo (from the repo root):

pip install -e .

• Install LeGrad (if not already installed):

pip install legrad

Make sure you have a working CUDA‑enabled PyTorch environment.

2. Open the notebook

From the repo root:

cd xai/perception_models
jupyter lab legrad_perception_encoder.ipynb

3. What the notebook does

The notebook shows how to:

1. Load a PE CoCa‑style vision encoder:
   • Uses pe.CLIP.from_config("PE-Core-B16-224", pretrained=True) and moves the model to CUDA.
2. Wrap the model with LeGrad:
   • LeWrapper lives in core/legrad_pe.py.
   • It hooks PE residual blocks and attention pooling so gradients can be used to build visual explanations.
3. Prepare inputs:
   • Build an image transform with transforms.get_image_transform(model.image_size).
   • Tokenize text prompts with transforms.get_text_tokenizer(model.context_length).
4. Run LeGrad:
   • Multi‑layer explanation:
     • heatmap = wrapped_model.compute_legrad_coca(text_emb, image=image_tensor)
   • Single‑layer explanation:
     • heatmap = wrapped_model.compute_legrad_coca_one_layer(text_emb, image=image_tensor, layer_idx=-1)
5. Visualize:
   • Convert the heatmap to numpy and use legrad.visualize (or standard plotting) to overlay it on the image.

4. Minimal code sketch (inside the notebook)

The core usage pattern is:

import core.vision_encoder.pe as pe
import core.vision_encoder.transforms as transforms
from core.legrad_pe import LeWrapper

model = pe.CLIP.from_config("PE-Core-B16-224", pretrained=True).cuda()
preprocess = transforms.get_image_transform(model.image_size)
tokenizer = transforms.get_text_tokenizer(model.context_length)

wrapped_model = LeWrapper(model, layer_index=-2)

You can then:

• Preprocess an input image with preprocess,
• Tokenize prompts with tokenizer,
• Encode text/image, and
• Call one of the compute_legrad_* methods to obtain a heatmap for visualization.