Coverage for transformer_lens/benchmarks/vision.py: 23%
94 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +0000
1"""Vision benchmarks for TransformerBridge (Phase 9).
3Tests that vision encoder models (ViT, DeiT) correctly handle pixel inputs
4through forward(), run_with_cache(), and produce stable representations —
5the hook/cache coverage that Phase 1 (HF parity on one forward) doesn't give
6non-text models. The audio analog is Phase 8 (audio.py).
7"""
9from typing import List, Optional
11import torch
13from transformer_lens.benchmarks.encoder_common import (
14 benchmark_encoder_cache,
15 benchmark_encoder_forward,
16 benchmark_encoder_representation_stability,
17)
18from transformer_lens.benchmarks.utils import (
19 BenchmarkResult,
20 BenchmarkSeverity,
21 build_modality_input,
22 is_tiny_test_model,
23)
24from transformer_lens.model_bridge import TransformerBridge
26# Top-level component-mapping names whose hook_out must be cached when the
27# architecture declares them (unembed = classifier head, absent on bare encoders).
28_CRITICAL_VISION_COMPONENTS = ("embed", "ln_final", "unembed")
31def benchmark_vision_forward(
32 bridge: TransformerBridge,
33 test_pixels: torch.Tensor,
34 reference_model: Optional[torch.nn.Module] = None,
35) -> BenchmarkResult:
36 """Benchmark forward pass with pixel input.
38 Compares bridge output against the HF native model on the same pixels.
39 Bare encoders (ViTModel) compare last_hidden_state; classification heads
40 compare logits.
42 Args:
43 bridge: TransformerBridge model to test
44 test_pixels: Pixel tensor [batch, channels, height, width]
45 reference_model: Optional HF reference model for comparison
46 """
47 return benchmark_encoder_forward(
48 bridge,
49 test_pixels,
50 name="vision_forward",
51 ref_input_key="pixel_values",
52 reference_model=reference_model,
53 )
56def benchmark_vision_cache(
57 bridge: TransformerBridge,
58 test_pixels: torch.Tensor,
59) -> BenchmarkResult:
60 """Benchmark run_with_cache() for vision models.
62 Verifies that critical vision hooks fire and produce valid tensors: the
63 patch embeddings, final layernorm, classifier head (when present), and the
64 first and last block.
66 Args:
67 bridge: TransformerBridge model to test
68 test_pixels: Pixel tensor [batch, channels, height, width]
69 """
70 return benchmark_encoder_cache(
71 bridge,
72 test_pixels,
73 name="vision_cache",
74 critical_components=_CRITICAL_VISION_COMPONENTS,
75 )
78def benchmark_vision_representation_stability(
79 bridge: TransformerBridge,
80 test_pixels: torch.Tensor,
81) -> BenchmarkResult:
82 """Benchmark representation stability under small pixel perturbations.
84 Args:
85 bridge: TransformerBridge model to test
86 test_pixels: Pixel tensor [batch, channels, height, width]
87 """
88 return benchmark_encoder_representation_stability(
89 bridge,
90 test_pixels,
91 name="vision_representation_stability",
92 )
95def benchmark_vision_embeddings(
96 bridge: TransformerBridge,
97 test_pixels: torch.Tensor,
98) -> BenchmarkResult:
99 """Verify patch-embedding hook outputs.
101 Checks that embed.hook_out produces [batch, seq, d_model] tensors with
102 non-degenerate values — the vision analog of the audio feature-extractor
103 check. Seq length is architecture-dependent (ViT: patches + CLS; DeiT:
104 patches + CLS + distillation token), so only lower-bounded here.
106 Args:
107 bridge: TransformerBridge model to test
108 test_pixels: Pixel tensor [batch, channels, height, width]
109 """
110 try:
111 if "embed" not in (bridge.adapter.component_mapping or {}):
112 return BenchmarkResult(
113 name="vision_embeddings",
114 severity=BenchmarkSeverity.SKIPPED,
115 message="Skipped: architecture declares no embed component",
116 )
118 with torch.no_grad():
119 _, cache = bridge.run_with_cache(test_pixels)
121 hook_key = "embed.hook_out"
122 if hook_key not in cache:
123 return BenchmarkResult(
124 name="vision_embeddings",
125 severity=BenchmarkSeverity.DANGER,
126 message=f"Hook '{hook_key}' not found in cache",
127 passed=False,
128 )
130 embeddings = cache[hook_key]
132 if embeddings.dim() != 3:
133 return BenchmarkResult(
134 name="vision_embeddings",
135 severity=BenchmarkSeverity.DANGER,
136 message=f"Expected 3D tensor [batch, seq, d_model], got {embeddings.dim()}D",
137 passed=False,
138 details={"shape": str(embeddings.shape)},
139 )
141 d_model = bridge.cfg.d_model
142 if embeddings.shape[0] != test_pixels.shape[0] or embeddings.shape[-1] != d_model:
143 return BenchmarkResult(
144 name="vision_embeddings",
145 severity=BenchmarkSeverity.DANGER,
146 message=f"Embedding shape {tuple(embeddings.shape)} does not match "
147 f"[batch={test_pixels.shape[0]}, seq, d_model={d_model}]",
148 passed=False,
149 details={"shape": str(embeddings.shape), "d_model": d_model},
150 )
152 is_all_zeros = embeddings.abs().max().item() == 0
153 has_nan = torch.isnan(embeddings).any().item()
154 has_inf = torch.isinf(embeddings).any().item()
156 if is_all_zeros or has_nan or has_inf:
157 issues = []
158 if is_all_zeros:
159 issues.append("all zeros")
160 if has_nan:
161 issues.append("NaN")
162 if has_inf:
163 issues.append("Inf")
164 return BenchmarkResult(
165 name="vision_embeddings",
166 severity=BenchmarkSeverity.DANGER,
167 message=f"Degenerate embedding values: {', '.join(issues)}",
168 passed=False,
169 details={"shape": str(embeddings.shape), "issues": issues},
170 )
172 return BenchmarkResult(
173 name="vision_embeddings",
174 severity=BenchmarkSeverity.INFO,
175 message=f"Patch embeddings OK: shape={embeddings.shape}, "
176 f"mean={embeddings.mean().item():.4f}, std={embeddings.std().item():.4f}",
177 details={
178 "shape": str(embeddings.shape),
179 "mean": embeddings.mean().item(),
180 "std": embeddings.std().item(),
181 },
182 )
184 except Exception as e:
185 return BenchmarkResult(
186 name="vision_embeddings",
187 severity=BenchmarkSeverity.ERROR,
188 message=f"Patch embedding check failed: {str(e)}",
189 passed=False,
190 )
193def benchmark_vision_classification_decode(
194 bridge: TransformerBridge,
195) -> BenchmarkResult:
196 """Benchmark image-classification decoding on a real image.
198 Loads the cats-image fixture, preprocesses it with the bridge's image
199 processor, and reports the top predicted labels — the vision analog of the
200 audio CTC decode, exercising the processor wiring that synthetic pixel
201 tensors don't. Skipped for bare encoders (no classifier head), tiny-random
202 models, and when no processor/datasets are available.
204 Args:
205 bridge: TransformerBridge model to test
206 """
207 model_name = getattr(bridge.cfg, "model_name", "")
208 if is_tiny_test_model(model_name):
209 return BenchmarkResult(
210 name="vision_classification_decode",
211 severity=BenchmarkSeverity.SKIPPED,
212 message="Skipped for tiny-random model (untrained classifier head)",
213 )
215 # Bare encoders return hidden states from return_type="logits", so gate on
216 # the adapter-declared classifier head rather than the output shape.
217 if "unembed" not in (bridge.adapter.component_mapping or {}):
218 return BenchmarkResult(
219 name="vision_classification_decode",
220 severity=BenchmarkSeverity.SKIPPED,
221 message="Skipped: bare encoder (no classifier head)",
222 )
224 processor = getattr(bridge, "processor", None)
225 if processor is None: 225 ↛ 232line 225 didn't jump to line 232 because the condition on line 225 was always true
226 return BenchmarkResult(
227 name="vision_classification_decode",
228 severity=BenchmarkSeverity.SKIPPED,
229 message="Skipped: no image processor available on the bridge",
230 )
232 try:
233 from datasets import load_dataset
235 ds = load_dataset("huggingface/cats-image", split="test", trust_remote_code=True)
236 image = ds[0]["image"]
238 inputs = processor(images=image, return_tensors="pt")
239 pixel_values = inputs["pixel_values"].to(bridge.cfg.device)
241 with torch.no_grad():
242 # Classifier heads return the logits tensor directly; bare encoders
243 # fall through to a BaseModelOutput with no logits attribute.
244 output = bridge(pixel_values, return_type="logits")
246 logits = output if isinstance(output, torch.Tensor) else getattr(output, "logits", None)
247 if logits is None or logits.ndim != 2:
248 return BenchmarkResult(
249 name="vision_classification_decode",
250 severity=BenchmarkSeverity.DANGER,
251 message="Classifier model produced no [batch, num_labels] logits",
252 passed=False,
253 )
255 k = min(5, logits.shape[-1])
256 top = torch.topk(logits[0], k=k)
257 id2label = getattr(getattr(bridge, "original_model", None), "config", None)
258 id2label = getattr(id2label, "id2label", None) or {}
259 top_labels = [id2label.get(int(i), str(int(i))) for i in top.indices]
261 return BenchmarkResult(
262 name="vision_classification_decode",
263 severity=BenchmarkSeverity.INFO,
264 message=f"Classification decode successful: top-1 = {top_labels[0]!r}",
265 details={
266 "top_labels": top_labels,
267 "top_logits": [round(float(v), 4) for v in top.values],
268 "logits_shape": str(logits.shape),
269 },
270 )
272 except ImportError:
273 return BenchmarkResult(
274 name="vision_classification_decode",
275 severity=BenchmarkSeverity.SKIPPED,
276 message="Skipped: 'datasets' package not available",
277 )
278 except Exception as e:
279 return BenchmarkResult(
280 name="vision_classification_decode",
281 severity=BenchmarkSeverity.ERROR,
282 message=f"Classification decode failed: {str(e)}",
283 passed=False,
284 )
287def run_vision_benchmarks(
288 bridge: TransformerBridge,
289 test_pixels: Optional[torch.Tensor] = None,
290 verbose: bool = True,
291) -> List[BenchmarkResult]:
292 """Run all vision benchmarks.
294 Args:
295 bridge: TransformerBridge model to test
296 test_pixels: Optional pixel tensor. If None, generates a synthetic input
297 shaped for this architecture from the HF config.
298 verbose: Whether to print progress
300 Returns:
301 List of BenchmarkResult objects
302 """
303 if test_pixels is None:
304 test_pixels = build_modality_input(bridge, device=bridge.cfg.device, dtype=bridge.cfg.dtype)
305 if test_pixels is None:
306 return [
307 BenchmarkResult(
308 name="vision_forward",
309 severity=BenchmarkSeverity.ERROR,
310 message="Could not build a pixel input for this model",
311 passed=False,
312 )
313 ]
315 results = []
317 if verbose:
318 print("1. Vision Forward Pass")
319 results.append(benchmark_vision_forward(bridge, test_pixels))
321 if verbose:
322 print("2. Vision Cache Verification")
323 results.append(benchmark_vision_cache(bridge, test_pixels))
325 if verbose:
326 print("3. Representation Stability")
327 results.append(benchmark_vision_representation_stability(bridge, test_pixels))
329 if verbose:
330 print("4. Patch Embedding Verification")
331 results.append(benchmark_vision_embeddings(bridge, test_pixels))
333 if verbose:
334 print("5. Classification Decoding")
335 results.append(benchmark_vision_classification_decode(bridge))
337 return results