feat(benchmarks): add evaluation framework for DINO-based compressors

mini-nav/compressors/__init__.py

@@ -1,5 +1,6 @@
+from .dino_compressor import DinoCompressor
 from .float_compressor import FloatCompressor
 from .int_compressor import IntCompressor
 from .train import train
 
-__all__ = ["train", "FloatCompressor", "IntCompressor"]
+__all__ = ["train", "FloatCompressor", "IntCompressor", "DinoCompressor"]

mini-nav/compressors/dino_compressor.py

@@ -0,0 +1,29 @@
+from typing import Optional, cast
+
+import torch.nn.functional as F
+from torch import nn
+from transformers import AutoModel, Dinov2Model
+
+
+class DinoCompressor(nn.Module):
+    def __init__(self, compressor: Optional[nn.Module] = None):
+        super().__init__()
+
+        self.dino = cast(
+            Dinov2Model,
+            AutoModel.from_pretrained("facebook/dinov2-large"),
+        )
+
+        self.compressor = compressor
+
+    def forward(self, inputs):
+        teacher_tokens = self.dino(**inputs).last_hidden_state  # [B,N,1024]
+
+        teacher_embed = teacher_tokens.mean(dim=1)
+        teacher_embed = F.normalize(teacher_embed, dim=-1)  # [B,1024]
+
+        if self.compressor is None:
+            return teacher_embed
+
+        feats, recon = self.compressor(teacher_tokens)
+        return feats

mini-nav/compressors/train.py

@@ -93,12 +93,10 @@ def train(
                 with torch.no_grad():
                     inputs = processor(imgs, return_tensors="pt").to(device)
 
-                    teacher_tokens = dino(**inputs).last_hidden_state
-                    # [B,N,1024]
+                    teacher_tokens = dino(**inputs).last_hidden_state  # [B,N,1024]
 
                     teacher_embed = teacher_tokens.mean(dim=1)
-                    teacher_embed = F.normalize(teacher_embed, dim=-1)
-                    # [B,1024]
+                    teacher_embed = F.normalize(teacher_embed, dim=-1)  # [B,1024]
 
                 # ---- student forward ----
                 z512, recon = compressor(teacher_tokens)