vision: DINOv3Stem, a timm DINOv3 ViT image stem (patch tokens, any aspect ratio)

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.
 

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