From f9cf65e172d693c663a134071b8d37a2e5bc3991 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 22:18:20 +0000 Subject: [PATCH] vision: DINOv3Stem, a timm DINOv3 ViT image stem (patch tokens, any aspect ratio) Same call convention as ResNet. Loaded through timm: the locked transformers (4.53.2) has no DINOv3 class, and the timm hub repos are ungated. RoPE-only, so image_size may be [H, W]; exposes grid_size. pretrained_reference_state_dict reads the checkpoint back from the local HF cache (never the network), so the pretrained-weight hash check covers this tower as it does the Qwen stems; an uncached checkpoint or a random tower is skipped. Against the real timm/vit_base_patch16_dinov3.lvd1689m file the key sets match (162 tensors) and every value is equal. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_019DbPWJCKtVg5c2PNHs7egw --- egomimic/models/hpt_nets.py | 118 +++++++++++++++++++++++++++++++++ tests/unit/test_dinov3_stem.py | 101 ++++++++++++++++++++++++++++ 2 files changed, 219 insertions(+) create mode 100644 tests/unit/test_dinov3_stem.py diff --git a/egomimic/models/hpt_nets.py b/egomimic/models/hpt_nets.py index 746dcf650e..4fe4ac0cac 100644 --- a/egomimic/models/hpt_nets.py +++ b/egomimic/models/hpt_nets.py @@ -956,6 +956,124 @@ def _local_snapshot_if_offline(model_name: str) -> str: return snapshot_download(model_name, local_files_only=True) +class DINOv3Stem(PretrainedWeights, PolicyStem): + """Pretrained DINOv3 ViT (timm), a drop-in replacement for ``ResNet``: + ``(B, T, N, 3, H, W)`` images in [0, 1] in, ``(B, patches, output_dim)`` out. + The stem owns the resize and the checkpoint's mean / std, so the config + must NOT normalize these images in its augs. + + Loaded through timm rather than transformers: the ``timm/*_dinov3`` hub + repos are not gated (``facebook/dinov3-*`` is), and DINOv3 in transformers + needs >= 4.56, above our pin. + + Args: + model_name: timm model id, e.g. ``vit_base_patch16_dinov3.lvd1689m`` + (86M) or ``vit_large_patch16_dinov3.lvd1689m`` (300M). + output_dim: width of the projected tokens. + freeze_backbone: freeze the tower, keep it in eval mode and run it + under ``no_grad``. + image_size: ``S`` or ``[H, W]``, multiples of the patch size. DINOv3 is + RoPE-based, so non-square inputs need no position interpolation. + pretrained: ``False`` builds a randomly initialised tower (tests). + tower_kwargs: timm architecture overrides (tests). + """ + + DEFAULT_MODEL = "vit_base_patch16_dinov3.lvd1689m" + _hpt_pretrained_attrs = ("tower",) + + def __init__( + self, + model_name: str = DEFAULT_MODEL, + output_dim: int = 10, + freeze_backbone: bool = True, + image_size=256, + pretrained: bool = True, + tower_kwargs: Optional[dict] = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + import timm + + self.model_name = model_name + self.pretrained = pretrained + self.tower = timm.create_model( + model_name, pretrained=pretrained, num_classes=0, **dict(tower_kwargs or {}) + ) + self.hidden_size = int(self.tower.embed_dim) + self.num_prefix_tokens = int(self.tower.num_prefix_tokens) + size = [image_size] * 2 if isinstance(image_size, int) else list(image_size) + patch = self.tower.patch_embed.patch_size[0] + if any(s % patch for s in size): + raise ValueError( + f"DINOv3Stem image_size {size} is not a multiple of {patch}" + ) + self.image_size = tuple(int(s) for s in size) + self.grid_size = tuple(s // patch for s in self.image_size) + cfg = self.tower.pretrained_cfg + self.register_buffer( + "image_mean", torch.tensor(cfg["mean"]).view(1, 3, 1, 1), persistent=False + ) + self.register_buffer( + "image_std", torch.tensor(cfg["std"]).view(1, 3, 1, 1), persistent=False + ) + + self.out_dim = output_dim + self.proj = nn.Linear(self.hidden_size, output_dim) + + self.freeze_backbone = freeze_backbone + if freeze_backbone: + for param in self.tower.parameters(): + param.requires_grad = False + + def train(self, mode: bool = True): + """Keep a frozen tower in eval mode regardless of the outer flag.""" + super().train(mode) + if self.freeze_backbone: + self.tower.eval() + return self + + def backbone_parameters(self): + """Pretrained-tower parameters (``proj`` stays in the main group).""" + return list(self.tower.parameters()) + + def pretrained_reference_state_dict(self) -> Optional[dict]: + """The cached timm checkpoint, keyed ``tower.*``. ``None`` for a random + tower, or when the file is not in the local HF cache.""" + if not self.pretrained: + return None + from huggingface_hub import hf_hub_download + from safetensors.torch import load_file + + try: + path = hf_hub_download( + self.tower.pretrained_cfg["hf_hub_id"], + "model.safetensors", + local_files_only=True, + ) + except Exception: # not cached, or not an HF-hosted timm model + return None + return {f"tower.{k}": v for k, v in load_file(path).items()} + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, *_, H, W = x.shape + x = x.reshape(-1, 3, H, W) + x = F.interpolate( + x, + size=self.image_size, + mode="bilinear", + align_corners=False, + antialias=True, + ) + x = (x - self.image_mean) / self.image_std + with torch.set_grad_enabled( + not self.freeze_backbone and torch.is_grad_enabled() + ): + feat = self.tower.forward_features(x) + # drop CLS + register tokens; RDT conditions on patch tokens only + feat = feat[:, self.num_prefix_tokens :] + return self.proj(feat.reshape(B, -1, self.hidden_size)) + + class _Qwen3BaseEncoder(PretrainedWeights, PolicyStem): """Shared base for Qwen3-Embedding stems used by HPT. diff --git a/tests/unit/test_dinov3_stem.py b/tests/unit/test_dinov3_stem.py new file mode 100644 index 0000000000..dfc1608478 --- /dev/null +++ b/tests/unit/test_dinov3_stem.py @@ -0,0 +1,101 @@ +"""``DINOv3Stem`` on a random-init tiny tower (no network, no weights).""" + +from __future__ import annotations + +import pytest +import torch + +from egomimic.models.hpt_nets import DINOv3Stem + +D, B = 32, 3 +TINY_TOWER = {"embed_dim": 32, "depth": 1, "num_heads": 2} + + +def _dino(**kwargs): + torch.manual_seed(0) + kwargs.setdefault("freeze_backbone", True) + return DINOv3Stem( + model_name="vit_small_patch16_dinov3", + output_dim=D, + image_size=[32, 48], + pretrained=False, + tower_kwargs=TINY_TOWER, + **kwargs, + ) + + +def test_dinov3_stem_returns_patch_tokens_only(): + stem = _dino() + assert stem.grid_size == (2, 3) + assert stem(torch.rand(B, 1, 1, 3, 90, 160)).shape == (B, 2 * 3, D) + # 2 frames on the time axis: past frame's patches first + assert stem(torch.rand(B, 2, 1, 3, 90, 160)).shape == (B, 2 * 2 * 3, D) + + +def test_dinov3_stem_frozen_tower(): + stem = _dino().train() + assert not stem.tower.training + stem(torch.rand(B, 1, 1, 3, 32, 48)).sum().backward() + assert all(p.grad is None for p in stem.tower.parameters()) + assert stem.proj.weight.grad is not None + assert stem.pretrained_submodules() == [stem.tower] + + +def test_dinov3_stem_finetuned_tower_gets_gradients(): + stem = _dino(freeze_backbone=False).train() + stem(torch.rand(B, 1, 1, 3, 32, 48)).sum().backward() + assert any(p.grad is not None for p in stem.tower.parameters()) + assert {id(p) for p in stem.backbone_parameters()} == { + id(p) for p in stem.tower.parameters() + } + + +def test_dinov3_stem_rejects_off_grid_size(): + with pytest.raises(ValueError, match="multiple"): + DINOv3Stem( + model_name="vit_small_patch16_dinov3", + image_size=250, + pretrained=False, + tower_kwargs=TINY_TOWER, + ) + + +def test_a_random_tower_is_not_verified(): + assert _dino().pretrained_reference_state_dict() is None + + +def test_the_verifier_checks_the_tower_against_its_cached_checkpoint( + tmp_path, monkeypatch +): + """The stem reads its timm checkpoint back from the HF cache, so #645's + hash check covers it like the Qwen stems.""" + import huggingface_hub + from safetensors.torch import save_file + + from egomimic.models.hpt_nets import verify_pretrained_weights + + stem = _dino() + path = tmp_path / "model.safetensors" + save_file({k: v.contiguous() for k, v in stem.tower.state_dict().items()}, path) + monkeypatch.setattr(huggingface_hub, "hf_hub_download", lambda *a, **k: str(path)) + stem.pretrained = True + + assert list(verify_pretrained_weights(stem, verbose=False)) == ["DINOv3Stem"] + with torch.no_grad(): + next(stem.tower.parameters()).add_(1.0) + with pytest.raises(RuntimeError, match="do not"): + verify_pretrained_weights(stem, verbose=False) + + +def test_an_uncached_checkpoint_is_skipped_not_fetched(monkeypatch): + import huggingface_hub + from huggingface_hub.errors import LocalEntryNotFoundError + + def offline(*a, **k): + assert k.get("local_files_only") is True + raise LocalEntryNotFoundError("not cached") + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", offline) + stem = _dino() + stem.pretrained = True + assert stem.pretrained_reference_state_dict() is None