feat(models): dinov2

.pre-commit-config.yaml

@@ -61,6 +61,7 @@ repos:
           - numpy
           - pandas-stubs
           - types-Pillow
+          - types-requests
           - types-toml
   - repo: https://github.com/codespell-project/codespell
     rev: v2.2.6

docs/source/pretrained/dino.py

@@ -0,0 +1,62 @@
+import torch
+import torchvision.transforms.v2 as tf_v2
+from einops.layers.torch import Rearrange
+
+from todd.datasets import IMAGENET_MEAN, IMAGENET_STD, COCODataset
+from todd.models.modules import DINO, DINOv2
+from todd.utils import get_image
+
+url = COCODataset.url('val', 2017, 39769)
+image = get_image(url)
+
+transforms = tf_v2.Compose([
+    torch.from_numpy,
+    Rearrange('h w c -> 1 c h w'),
+    tf_v2.Resize(256, tf_v2.InterpolationMode.BICUBIC),
+    tf_v2.CenterCrop(256),
+    tf_v2.ToDtype(torch.float32, scale=True),
+    tf_v2.Normalize(IMAGENET_MEAN, IMAGENET_STD),
+])
+tensor = transforms(image)
+
+dino = DINO()
+dino.load_pretrained('pretrained/dino/dino_vitbase16_pretrain.pth')
+cls_, x = dino(tensor, False)
+assert torch.allclose(
+    cls_[:, :3],
+    torch.tensor([[2.3480, -5.4728, 3.2335]]),
+    atol=1e-4,
+)
+assert torch.allclose(
+    x[:, :3, :3],
+    torch.tensor([[
+        [-1.5503, -2.4765, 1.1179],
+        [-2.6262, -1.3692, -0.1952],
+        [-2.4387, 0.5866, 0.0967],
+    ]]),
+    atol=1e-4,
+)
+
+dinov2 = DINOv2(
+    patch_size=14,
+    patch_wh=(37, 37),
+    width=1024,
+    depth=24,
+    num_heads=16,
+)
+dinov2.load_pretrained('pretrained/dino/dinov2_vitl14_pretrain.pth')
+cls_, x = dinov2(tensor, False)
+assert torch.allclose(
+    cls_[:, :3],
+    torch.tensor([[-2.7620, -2.1128, 0.4863]]),
+    atol=1e-4,
+)
+assert torch.allclose(
+    x[:, :3, :3],
+    torch.tensor([[
+        [2.4025, -1.0418, 1.6840],
+        [1.6887, 0.0430, 3.5852],
+        [0.5082, -0.5240, 0.8340],
+    ]]),
+    atol=1e-4,
+)

docs/source/pretrained/dino.rst

@@ -0,0 +1,25 @@
+DINO
+====
+
+.. code-block:: bash
+
+    root=pretrained/dino
+    mkdir -p ${root}
+
+DINO
+----
+
+https://github.com/facebookresearch/dino
+
+.. code-block:: bash
+
+    wget https://dl.fbaipublicfiles.com/dino/dino_vitbase16_pretrain/dino_vitbase16_pretrain.pth -P ${root}
+
+DINOv2
+------
+
+https://github.com/facebookresearch/dinov2
+
+.. code-block:: bash
+
+    wget https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_pretrain.pth -P ${root}

pyproject.toml

@@ -79,6 +79,7 @@ lint = [
     'pylint',
     'types-Pillow',
     'types-pycocotools',
+    'types-requests',
     'types-toml',
     'types-tqdm',
 ]

tests/test_patches/test_torch/test_torch_builtins.py → tests/test_patches/test_torch/test_aten.py

@@ -1,4 +1,4 @@
-from todd.patches.torch.builtins import random_int
+from todd.patches.torch.aten import random_int
 from todd.utils import init_seed
 
 

tests/test_tasks/test_odkd/test_distillers/test_fgd.py

@@ -1,7 +1,6 @@
 import os
 import pathlib
 
-import mmcv.cnn
 import pytest
 import torch
 
@@ -17,11 +16,13 @@
 class TestFGD:
 
     @pytest.fixture(autouse=True)
-    def mmcv(self) -> None:
+    def setup(self) -> None:
+        from mmcv.cnn import ContextBlock
+
         kd.distillers.KDAdaptRegistry.register_(
             'mmcv_ContextBlock',
             force=True,
-        )(mmcv.cnn.ContextBlock)
+        )(ContextBlock)
 
     # @pytest.fixture
     # def distiller(self, data_dir: pathlib.Path) -> BaseDistiller:

todd/datasets/coco.py

@@ -32,6 +32,7 @@
     from todd.tasks.object_detection import BBox, FlattenBBoxesXYWH
 
 Split = Literal['train', 'val']
+Year = Literal[2014, 2017]
 
 
 class BaseKeys(KeysProtocol[str], ABC):
@@ -173,11 +174,15 @@ class COCODataset(BaseDataset[COCO, T]):
     DATA_ROOT = pathlib.Path('data/coco')
     ANNOTATIONS_ROOT = DATA_ROOT / 'annotations'
 
+    @classmethod
+    def url(cls, split: Split, year: Year, id_: int) -> str:
+        return f'http://images.cocodataset.org/{split}{year}/{id_:012d}.jpg'
+
     def __init__(
         self,
         *args,
         split: Split,
-        year: Literal[2014, 2017] = 2017,
+        year: Year = 2017,
         access_layer: PILAccessLayer | None = None,
         annotations_file: pathlib.Path | str | None = None,
         **kwargs,

todd/models/__init__.py

@@ -1,5 +1,4 @@
-from . import filters, losses, shadows, utils
-from .dino import *
+from . import filters, losses, modules, shadows, utils
 from .frozens import *
 from .mean_std import *
 from .norms import *

todd/models/modules/__init__.py

@@ -0,0 +1,2 @@
+from .dino import *
+from .vit import *