| from pathlib import Path |
| import os |
| from io import BytesIO |
| from urllib.request import urlopen |
| import torch |
| from PIL import Image |
| from transformers import AutoImageProcessor, ViTForImageClassification |
| from vlmprobe import Prober |
| from vlmprobe import visualization as viz |
| import matplotlib.pyplot as plt |
|
|
| |
| |
| model_id = os.environ.get("VLM_PROBING_MODEL", "google/vit-base-patch16-224") |
| device = os.environ.get("VLM_PROBING_DEVICE", "cuda:0" if torch.cuda.is_available() else "cpu") |
| model = ViTForImageClassification.from_pretrained(model_id, attn_implementation="eager").to(device).eval() |
| processor = AutoImageProcessor.from_pretrained(os.environ.get("VLM_PROBING_PROCESSOR", model_id)) |
| probe = Prober(model, processor) |
|
|
| |
| image_path = os.environ.get("VLM_PROBING_IMAGE") |
| image = (Image.open(image_path) if image_path else Image.open(BytesIO(urlopen( |
| "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/coco_sample.png" |
| ).read()))).convert("RGB") |
| inputs = {k: v.to(device) for k, v in processor(images=image, return_tensors="pt").items()} |
|
|
| |
| mean = torch.tensor(processor.image_mean)[:, None, None] |
| std = torch.tensor(processor.image_std)[:, None, None] |
| rgb = (inputs["pixel_values"][0].detach().cpu() * std + mean).clamp(0, 1) |
| processed_image = Image.fromarray((rgb.permute(1, 2, 0).numpy() * 255).astype("uint8")) |
| patch = model.config.patch_size |
| grid_shape = tuple(int(n // patch) for n in inputs["pixel_values"].shape[-2:]) |
| visual_positions = list(range(1, 1 + grid_shape[0] * grid_shape[1])) |
| last = probe.describe()["methods"]["attention.profile"]["sites"][-1] |
| output_dir = Path(os.environ.get("VLM_PROBING_OUTPUT", "outputs/attention-tutorial")) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| result = probe.attention.profile( |
| layers=[last], queries=[0], groups={"patches": "visual", "cls": [0]}, |
| include_attention=True, |
| ).run(inputs) |
| print("Entropy [L,B,H,Q]:", result.tensors["entropy"].shape) |
| print("Selected query entropy:", result.tensors["entropy"][0, 0, :, 0]) |
| print("Group order:", result.metadata["layer_results"][0]["groups"]) |
| ax = viz.plot_attention_overlay( |
| processed_image, result, tensor="attention", grid_shape=grid_shape, |
| visual_positions=visual_positions, query_positions=[0], |
| signed=False, image_alpha=0.25, alpha=1.0, |
| title="Raw CLS attention (mean heads)", |
| ) |
| ax.figure.savefig(output_dir / "attention-profile.png", dpi=160, bbox_inches="tight") |
| result.save(output_dir / "attention-profile") |
| plt.show() |
|
|
