Upload 162 files

DeepForestVisionV2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21f12d6a6610efd8c69f04302b6dc7cdec3c8e75a8306e70348469af60c2c44e
+size 343226287

LICENSE.md

@@ -0,0 +1,66 @@
+# DINOv3 License
+
+*Last Updated: August 19, 2025*
+
+**“Agreement”** means the terms and conditions for use, reproduction, distribution and modification of the DINO Materials set forth herein.
+
+**“DINO Materials”** means, collectively, Documentation and the models, software and algorithms, including machine-learning model code, trained model weights, inference-enabling code, training-enabling code, fine-tuning enabling code, and other elements of the foregoing distributed by Meta and made available under this Agreement.
+
+**“Documentation”** means the specifications, manuals and documentation accompanying
+DINO Materials distributed by Meta.
+
+**“Licensee”** or **“you”** means you, or your employer or any other person or entity (if you are entering into this Agreement on such person or entity’s behalf), of the age required under applicable laws, rules or regulations to provide legal consent and that has legal authority to bind your employer or such other person or entity if you are entering in this Agreement on their behalf.
+
+**“Meta”** or **“we”** means Meta Platforms Ireland Limited (if you are located in or, if you are an entity, your principal place of business is in the EEA or Switzerland) or Meta Platforms, Inc. (if you are located outside of the EEA or Switzerland).
+
+**“Sanctions”** means any economic or trade sanctions or restrictions administered or enforced by the United States (including the Office of Foreign Assets Control of the U.S. Department of the Treasury (“OFAC”), the U.S. Department of State and the U.S. Department of Commerce), the United Nations, the European Union, or the United Kingdom.
+
+**“Trade Controls”** means any of the following: Sanctions and applicable export and import controls.
+
+By clicking “I Accept” below or by using or distributing any portion or element of the DINO Materials, you agree to be bound by this Agreement.
+
+## 1. License Rights and Redistribution.
+
+a. <ins>Grant of Rights</ins>. You are granted a non-exclusive, worldwide, non-transferable and royalty-free limited license under Meta’s intellectual property or other rights owned by Meta embodied in the DINO Materials to use, reproduce, distribute, copy, create derivative works of, and make modifications to the DINO Materials.
+
+b. <ins>Redistribution and Use</ins>.
+
+i. Distribution of DINO Materials, and any derivative works thereof, are subject to the terms of this Agreement. If you distribute or make the DINO Materials, or any derivative works thereof, available to a third party, you may only do so under the terms of this Agreement and you shall provide a copy of this Agreement with any such DINO Materials.
+
+ii.  If you submit for publication the results of research you perform on, using, or otherwise in connection with DINO Materials, you must acknowledge the use of DINO Materials in your publication.
+
+iii. Your use of the DINO Materials must comply with applicable laws and regulations, including Trade Control Laws and applicable privacy and data protection laws.
+
+iv. Your use of the DINO Materials will not involve or encourage others to reverse engineer, decompile or discover the underlying components of the DINO Materials.
+
+v. You are not the target of Trade Controls and your use of DINO Materials must comply with Trade Controls. You agree not to use, or permit others to use, DINO Materials for any activities subject to the International Traffic in Arms Regulations (ITAR) or end uses prohibited by Trade Controls, including those related to military or warfare purposes, nuclear industries or applications, espionage, or the development or use of guns or illegal weapons.
+
+## 2. User Support.
+
+Your use of the DINO Materials is done at your own discretion; Meta does not process any information nor provide any service in relation to such use.  Meta is under no obligation to provide any support services for the DINO Materials. Any support provided is “as is”, “with all faults”, and without warranty of any kind.
+
+## 3. Disclaimer of Warranty.
+
+UNLESS REQUIRED BY APPLICABLE LAW, THE DINO MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN “AS IS” BASIS, WITHOUT WARRANTIES OF ANY KIND, AND META DISCLAIMS ALL WARRANTIES OF ANY KIND, BOTH EXPRESS AND IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING OR REDISTRIBUTING THE DINO MATERIALS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE DINO MATERIALS AND ANY OUTPUT AND RESULTS.
+
+## 4. Limitation of Liability.
+
+IN NO EVENT WILL META OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY DIRECT OR INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF META OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
+
+## 5. Intellectual Property.
+
+a. Subject to Meta’s ownership of DINO Materials and derivatives made by or for Meta, with respect to any derivative works and modifications of the DINO Materials that are made by you, as between you and Meta, you are and will be the owner of such derivative works and modifications.
+
+b. If you institute litigation or other proceedings against Meta or any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the DINO Materials, outputs or results, or any portion of any of the foregoing, constitutes infringement of intellectual property or other rights owned or licensable by you, then any licenses granted to you under this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and hold harmless Meta from and against any claim by any third party arising out of or related to your use or distribution of the DINO Materials.
+
+## 6. Term and Termination.
+
+The term of this Agreement will commence upon your acceptance of this Agreement or access to the DINO Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. Meta may terminate this Agreement if you are in breach of any term or condition of this Agreement. Upon termination of this Agreement, you shall delete and cease use of the DINO Materials. Sections 3, 4 and 7 shall survive the termination of this Agreement.
+
+## 7. Governing Law and Jurisdiction.
+
+This Agreement will be governed and construed under the laws of the State of California without regard to choice of law principles, and the UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The courts of California shall have exclusive jurisdiction of any dispute arising out of this Agreement.
+
+## 8. Modifications and Amendments.
+
+Meta may modify this Agreement from time to time; provided that they are similar in spirit to the current version of the Agreement, but may differ in detail to address new problems or concerns. All such changes will be effective immediately. Your continued use of the DINO Materials after any modification to this Agreement constitutes your agreement to such modification. Except as provided in this Agreement, no modification or addition to any provision of this Agreement will be binding unless it is in writing and signed by an authorized representative of both you and Meta.

dinov3/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+__version__ = "0.0.1"

dinov3/checkpointer/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from .checkpointer import (
+    CheckpointRetentionPolicy,
+    cleanup_checkpoint,
+    find_all_checkpoints,
+    find_latest_checkpoint,
+    init_fsdp_model_from_checkpoint,
+    init_model_from_checkpoint_for_evals,
+    keep_checkpoint_copy,
+    keep_last_n_checkpoints,
+    load_checkpoint,
+    register_dont_save_hooks,
+    save_checkpoint,
+)

dinov3/checkpointer/checkpointer.py

@@ -0,0 +1,352 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+"""
+Suggested file structure:
+
+output_dir/
+|-- ckpt/
+|   |-- 0/
+|   |-- 99/
+|   |-- 199/
+|   |-- 199_keep/
+|   |-- 299/
+|   `-- ...
+`-- eval/
+    `-- 0/
+    `-- 99/
+        `-- ckpt/
+
+Distributed checkpointer docs:
+- https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html
+- https://pytorch.org/docs/stable/distributed.checkpoint.html
+"""
+
+import logging
+import shutil
+import subprocess
+import tempfile
+from enum import Enum
+from pathlib import Path
+from typing import List, Sequence, Set
+
+import torch
+import torch.distributed as dist
+import torch.distributed.checkpoint as dcp
+import torch.distributed.checkpoint.filesystem as dcpfs
+import torch.distributed.checkpoint.state_dict as dcpsd
+from torch.distributed.checkpoint.stateful import Stateful
+
+logger = logging.getLogger("dinov3")
+
+
+class CheckpointRetentionPolicy(Enum):
+    ALL = "all"  # keep all checkpoints
+    BEST = "best"
+    LAST = "last"
+    LAST_AND_BEST = "last_and_best"
+    NONE = "none"  # do not keep any checkpoints
+
+    @property
+    def keep_filters(self) -> Set[str]:
+        """Files that match these patterns are not deleted by cleanup"""
+        if self == CheckpointRetentionPolicy.LAST:
+            return set(["final"])
+        if self == CheckpointRetentionPolicy.BEST:
+            return set(["best"])
+        if self == CheckpointRetentionPolicy.LAST_AND_BEST:
+            return set(["final", "best"])
+        if self == CheckpointRetentionPolicy.ALL:
+            return set()
+        return set()
+
+    @property
+    def max_to_keep(self) -> int | None:
+        """
+        maximum "periodic" checkpoints to keep concurrently, ie. saved with `step` and not `save`. `None` for keep all
+        """
+        if self == CheckpointRetentionPolicy.ALL:
+            return None
+        return 1
+
+
+def save_checkpoint(
+    ckpt_dir: str | Path,  # output_dir/ckpt/199
+    *,
+    iteration: int | str,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer | None = None,
+    overwrite: bool = True,
+    process_group: dist.ProcessGroup = None,
+    **others: Stateful,
+):
+    """Save a plain/DDP/FSDP/FSDP2 model, its optimizer, an integer iteration and other stateful objects."""
+    rank = torch.distributed.get_rank(group=process_group)
+
+    # Rank 0 checks if the checkpoint directory exists, but all ranks need to know if if exists,
+    # so they can raise an error when overwrite is False. If overwrite is True, rank 0 will delete it
+    # and other ranks wait for the deletion to finish.
+    ckpt_dir = Path(ckpt_dir)
+    ckpt_dir_exists = [ckpt_dir.exists() if rank == 0 else None]
+    src_rank = 0
+    if process_group is not None:
+        src_rank = torch.distributed.get_global_rank(group=process_group, group_rank=0)
+    torch.distributed.broadcast_object_list(ckpt_dir_exists, src=src_rank, group=process_group)
+    ckpt_dir_exists = ckpt_dir_exists[0]
+    if ckpt_dir_exists:
+        if overwrite:
+            if rank == 0:
+                if ckpt_dir.is_dir():
+                    shutil.rmtree(ckpt_dir)
+                else:
+                    ckpt_dir.unlink()
+                logger.info(f"Deleted: {ckpt_dir}")
+            torch.distributed.barrier(group=process_group)
+        else:
+            raise RuntimeError(f"Checkpoint already exists: {ckpt_dir}")
+
+    # Rank 0 creates a temporary directory for the checkpoint and broadcasts the name to all ranks.
+    ckpt_dir.parent.mkdir(parents=True, exist_ok=True)
+    ckpt_dir_tmp = [tempfile.mkdtemp(dir=ckpt_dir.parent, prefix=ckpt_dir.name) if rank == 0 else None]
+    torch.distributed.broadcast_object_list(ckpt_dir_tmp, src=src_rank, group=process_group)
+    ckpt_dir_tmp = Path(ckpt_dir_tmp[0])
+
+    to_save = {"iteration": iteration}
+    to_save["model"] = dcpsd.get_model_state_dict(model)
+    if optimizer is not None:
+        to_save["optimizer"] = dcpsd.get_optimizer_state_dict(model, optimizer)
+    to_save.update(others)
+    dcp.save(
+        to_save,
+        storage_writer=dcpfs.FileSystemWriter(ckpt_dir_tmp),
+        process_group=process_group,
+    )
+
+    # Rank 0 renames the temporary directory to the final checkpoint directory. All ranks wait for the rename.
+    if rank == 0:
+        ckpt_dir_tmp.rename(ckpt_dir)
+    torch.distributed.barrier()
+
+    logger.info(f"Saved: {ckpt_dir}")
+
+
+def load_checkpoint(
+    ckpt_dir: str | Path,  # output_dir/ckpt/199
+    *,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer | None = None,
+    strict_loading: bool = True,
+    process_group: dist.ProcessGroup = None,
+    **others: Stateful,
+) -> int | None:
+    """
+    Load a plain/DDP/FSDP/FSDP2 model, its optimizer, an integer iteration and other stateful objects.
+    Can you take a checkpoint saved on N ranks and load it on M ranks? Sure you can!
+    Activation checkpointing and torch-compile can also be different between save and load, no problem.
+    """
+    ckpt_dir = Path(ckpt_dir)
+    to_load = {"iteration": None}
+    to_load["model"] = dcpsd.get_model_state_dict(model)
+    if optimizer is not None:
+        to_load["optimizer"] = dcpsd.get_optimizer_state_dict(model, optimizer)
+    to_load.update(others)
+    dcp.load(
+        to_load,
+        storage_reader=dcpfs.FileSystemReader(ckpt_dir),
+        planner=dcp.default_planner.DefaultLoadPlanner(allow_partial_load=not strict_loading),
+        process_group=process_group,
+    )
+    iteration = to_load["iteration"]
+    dcpsd.set_model_state_dict(model, to_load["model"])
+    if optimizer is not None:
+        dcpsd.set_optimizer_state_dict(model, optimizer, to_load["optimizer"])
+    logger.info(f"Loaded: {ckpt_dir}")
+    return iteration
+
+
+def register_dont_save_hooks(module: torch.nn.Module, dont_save: Sequence[str]):
+    """
+    Registers save/load state dict hooks such that the weights in `dont_save` are not persisted in the checkpoint.
+
+    Typical use case: a classification model composed of a frozen backbone and a trainable head.
+    If the frozen backbone is loaded from torch hub, it does't make sense to save a copy of it in each checkpoint.
+    """
+
+    def state_dict_post_hook(module, state_dict, prefix, local_metadata):
+        # Remove frozen weights so they won't get saved.
+        # If this module is not the top-level module, its weights will have a prefix in the state dict.
+        nonlocal _dont_save
+        for k in _dont_save:
+            del state_dict[prefix + k]
+
+    def load_state_dict_pre_hook(
+        module,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        # This pre hook exists only to pass the prefix to the post hook when loading the state dict.
+        nonlocal _prefix
+        assert _prefix is None
+        _prefix = prefix
+
+    def load_state_dict_post_hook(module, incompatible_keys):
+        # Remove the frozen weights from the missing keys so they don't raise an error.
+        nonlocal _prefix
+        assert _prefix is not None
+        to_remove = []
+        for missing_key in incompatible_keys.missing_keys:
+            k = missing_key.removeprefix(_prefix)
+            k = k.replace("_checkpoint_wrapped_module.", "")  # Added by activation checkpointing
+            if k in _dont_save:
+                to_remove.append(missing_key)
+        for r in to_remove:
+            incompatible_keys.missing_keys.remove(r)
+        _prefix = None
+
+    _dont_save = set(name.replace("_checkpoint_wrapped_module.", "") for name in dont_save)
+    _prefix = None
+    module.register_state_dict_post_hook(state_dict_post_hook)
+    module.register_load_state_dict_pre_hook(load_state_dict_pre_hook)
+    module.register_load_state_dict_post_hook(load_state_dict_post_hook)
+
+
+def find_all_checkpoints(ckpt_dir: Path | str) -> list[Path]:
+    """Find all checkpoints in a directory, i.e. subdirs with integer name. Sorted from first to last."""
+    ckpt_dir = Path(ckpt_dir)
+    if not ckpt_dir.is_dir():
+        return []
+    checkpoints = [p for p in ckpt_dir.iterdir() if p.is_dir() and _is_int(p.name)]
+    checkpoints.sort(key=lambda p: int(p.name))
+    return checkpoints
+
+
+def find_latest_checkpoint(ckpt_dir: Path | str) -> Path | None:
+    """Find the latest checkpoint in a directory, i.e. the subdir with the highest integer name."""
+    checkpoints = find_all_checkpoints(ckpt_dir)
+    if len(checkpoints) == 0:
+        return None
+    return checkpoints[-1]
+
+
+def keep_last_n_checkpoints(ckpt_dir: Path | str, n: int | None):
+    """In a directory with integer-named subdirs, keep only the n subdirs with the highest number."""
+    if n is None:
+        return
+    checkpoints = find_all_checkpoints(ckpt_dir)
+    for ckpt_dir in checkpoints[:-n]:
+        try:
+            shutil.rmtree(ckpt_dir)
+            logger.info(f"Deleted: {ckpt_dir}")
+        except Exception:
+            logger.exception(f"Failed to delete: {ckpt_dir}")
+
+
+def keep_checkpoint_copy(src: Path | str):
+    """Copy a file/directory next to itself with a _keep suffix. Files are hardlinked."""
+    src = Path(src)
+    dst = src.parent / f"{src.name}_keep"
+    subprocess.check_output(["cp", "--recursive", "--link", src, dst])
+    logger.info(f"Copied: {src} -> {dst}")
+
+
+def _is_int(s: str) -> bool:
+    try:
+        int(s)
+        return True
+    except ValueError:
+        return False
+
+
+# Initialize a FSDP2 model from DCP or PyTorch standard checkpoint
+def init_fsdp_model_from_checkpoint(
+    model: torch.nn.Module,
+    checkpoint_path: str,
+    skip_load_keys: List[str] | None = None,
+    keys_not_sharded: List[str] | None = None,
+    process_group: dist.ProcessGroup = None,
+):
+    if not Path(checkpoint_path).is_dir():  # PyTorch standard checkpoint
+        logger.info(f"Loading pretrained weights from {checkpoint_path}")
+        chkpt = torch.load(checkpoint_path, map_location="cpu")["teacher"]
+        from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
+
+        if process_group is None:
+            world_mesh = init_device_mesh(
+                "cuda",
+                mesh_shape=(dist.get_world_size(),),
+                mesh_dim_names=("dp",),
+            )
+        else:
+            world_mesh = DeviceMesh.from_group(process_group, "cuda")
+        chkpt = {
+            key: (
+                torch.distributed.tensor.distribute_tensor(tensor, world_mesh, src_data_rank=None)
+                if not any(key_not_sharded in key for key_not_sharded in keys_not_sharded)
+                else tensor
+            )
+            for key, tensor in chkpt.items()
+        }
+        model.load_state_dict(
+            {
+                key: tensor
+                for key, tensor in chkpt.items()
+                if not any(skip_load_key in key for skip_load_key in skip_load_keys)
+            }
+        )
+    else:  # DCP checkpoint
+        load_checkpoint(ckpt_dir=checkpoint_path, model=model, process_group=process_group)
+
+
+# Initialize a standard non distributed PyTorch model from PyTorch standard checkpoint for evals
+def init_model_from_checkpoint_for_evals(
+    model: torch.nn.Module, pretrained_weights: str | Path, checkpoint_key: str = None
+):
+    state_dict = torch.load(pretrained_weights, map_location="cpu")
+    if checkpoint_key is not None and checkpoint_key in state_dict:
+        logger.info(f"Take key {checkpoint_key} in provided checkpoint dict")
+        state_dict = state_dict[checkpoint_key]
+    # remove `module.` prefix
+    state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+    # remove `backbone.` prefix induced by multicrop wrapper
+    state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
+    msg = model.load_state_dict(state_dict, strict=False)
+    logger.info("Pretrained weights found at {} and loaded with msg: {}".format(pretrained_weights, msg))
+
+
+def cleanup_checkpoint(ckpt_dir: str, checkpoint_retention_policy: CheckpointRetentionPolicy):
+    """
+    ckpt_dir is the directory containing each individual checkpoint directories (either at iteration, best (validation performance) or final)
+    |-- ckpt_dir/
+    |   |-- 0/
+    |       |--checkpoint.pth  or dcp_sharded_checkpoint_dir
+    |   |-- 99/
+            |--checkpoint.pth or dcp_sharded_checkpoint_dir
+    |   |-- 199/
+            |--checkpoint.pth or dcp_sharded_checkpoint_dir
+    |   |-- best/
+            |--checkpoint.pth or dcp_sharded_checkpoint_dir
+    |   |-- 299/
+            |--checkpoint.pth or dcp_sharded_checkpoint_dir
+    |   |-- final/
+            |--checkpoint.pth or dcp_sharded_checkpoint_dir
+    """
+    ckpt_dir = Path(ckpt_dir)
+    if not ckpt_dir.is_dir():
+        return []
+    checkpoint_filters = checkpoint_retention_policy.keep_filters
+    checkpoints = [p for p in ckpt_dir.iterdir() if p.is_dir()]
+    for checkpoint in checkpoints:
+        if checkpoint in checkpoint_filters:
+            continue
+        try:
+            shutil.rmtree(checkpoint)
+            logger.info(f"Deleted: {checkpoint}")
+        except Exception:
+            logger.exception(f"Failed to delete: {checkpoint}")

dinov3/configs/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from .config import (
+    DinoV3SetupArgs,
+    apply_scaling_rules_to_cfg,
+    exit_job,
+    get_cfg_from_args,
+    get_default_config,
+    setup_config,
+    setup_job,
+    setup_multidistillation,
+    write_config,
+)

dinov3/configs/config.py

@@ -0,0 +1,222 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import logging
+import math
+import os
+import pathlib
+import sys
+from dataclasses import dataclass, field
+from datetime import timedelta
+from typing import Any, List, Optional, Sequence, Tuple
+
+from omegaconf import DictConfig, OmegaConf
+
+import dinov3.distributed as distributed
+from dinov3.logging import cleanup_logging, setup_logging
+from dinov3.utils import fix_random_seeds, get_conda_env, get_sha
+
+logger = logging.getLogger("dinov3")
+
+
+@dataclass
+class DinoV3SetupArgs:
+    config_file: str
+    pretrained_weights: str | None = None
+    shard_unsharded_model: bool = False
+    output_dir: str = ""
+    opts: List[Any] = field(default_factory=lambda: [])
+
+    def __post_init__(self):
+        # When loaded from benchmark.yaml, self.opts is a frozen omegaconf.ListConfig,
+        # which works everywhere except when we want to modify it or when
+        # we try to json-serialize it. So we convert it to a regular list here.
+        if OmegaConf.is_config(self.opts):
+            self.opts = OmegaConf.to_object(self.opts)
+
+
+def apply_scaling_rules_to_cfg(cfg):  # to fix
+    assert distributed.is_enabled(), "Setup distributed to get global size !"
+    if "schedules" in cfg:
+        # For schedules v2, the scaling rules are applied when building the schedules, the config is not modified
+        return cfg
+
+    if cfg.optim.scaling_rule == "linear_wrt_256":
+        old_lr = cfg.optim.lr
+        cfg.optim.lr *= cfg.train.batch_size_per_gpu * distributed.get_world_size() / 256.0
+        logger.info(f"linear scaling learning rate; old: {old_lr}, new: {cfg.optim.lr}")
+    elif cfg.optim.scaling_rule == "sqrt_wrt_1024":
+        old_lr = cfg.optim.lr
+        cfg.optim.lr *= 4 * math.sqrt(cfg.train.batch_size_per_gpu * distributed.get_world_size() / 1024.0)
+        logger.info(f"sqrt scaling learning rate; old: {old_lr}, new: {cfg.optim.lr}")
+    return cfg
+
+
+def write_config(cfg, output_dir, name="config.yaml"):
+    logger.info(OmegaConf.to_yaml(cfg))
+    output_dir = os.path.abspath(output_dir)
+    saved_cfg_path = os.path.join(output_dir, name)
+    with open(saved_cfg_path, "w") as f:
+        OmegaConf.save(config=cfg, f=f)
+    return saved_cfg_path
+
+
+def get_default_config() -> DictConfig:
+    p = pathlib.Path(__file__).parent / "ssl_default_config.yaml"
+    return OmegaConf.load(p)
+
+
+def get_cfg_from_args(args: DinoV3SetupArgs, multidistillation=False, strict=True):
+    overrides = [*args.opts]
+    if args.output_dir is not None:
+        overrides.append(f"train.output_dir={os.path.realpath(args.output_dir)}")
+
+    # Config file
+    cfg = OmegaConf.load(args.config_file)
+
+    # Command line overrides
+    opts_cfg = OmegaConf.from_cli(overrides)
+
+    if multidistillation:
+        cfg = OmegaConf.merge(cfg, opts_cfg)
+    else:
+        # Default config
+        default_cfg = get_default_config()
+        if strict:
+            OmegaConf.set_struct(default_cfg, True)
+        cfg = OmegaConf.merge(default_cfg, cfg, opts_cfg)
+    return cfg
+
+
+def setup_config(args: DinoV3SetupArgs, strict_cfg=True):
+    """
+    Create configs and perform basic setups.
+    """
+    # Create the cfg with OmegaConf
+    cfg = get_cfg_from_args(args, strict=strict_cfg)
+    # setup distributed, logging, and random seeds
+    logger.info("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
+    # dump config before modifying so it can be reloaded
+    if args.output_dir is not None:
+        write_config(cfg, args.output_dir)
+    # modify the config inplace by applying scaling rules
+    apply_scaling_rules_to_cfg(cfg)
+    return cfg
+
+
+def _enumerate_all_subgroup_ranks(all_subgroup_rank_spans: Sequence[Tuple[int, int]]):
+    """Expands a specification of process subgroups from spans to enumerated ranks.
+
+    Args:
+       all_group_rank_spans: a sequence of rank spans (first rank, last rank),
+           one for each process group. Example: ((0, 1), (2, 3), (4, 7)).
+    """
+    for first, last in all_subgroup_rank_spans:
+        assert first <= last
+    return tuple(tuple(range(first, last + 1)) for first, last in all_subgroup_rank_spans)
+
+
+def setup_multidistillation(args: DinoV3SetupArgs):
+    base_output_dir = args.output_dir
+    os.makedirs(args.output_dir, exist_ok=True)
+    # get config file for this rank
+    base_cfg = OmegaConf.load(args.config_file)
+    assert base_cfg.multidistillation.enabled
+
+    global_batch_size = base_cfg.multidistillation.global_batch_size
+
+    distributed.enable(overwrite=True)
+    seed = getattr(args, "seed", 0)
+    rank = distributed.get_rank()
+
+    # build process subgroups
+    all_subgroup_rank_spans = tuple(
+        (student.ranks_range[0], student.ranks_range[1] - 1) for student in base_cfg.multidistillation.students
+    )
+    all_subgroup_ranks = _enumerate_all_subgroup_ranks(all_subgroup_rank_spans)
+    distributed.new_subgroups(all_subgroup_ranks)
+
+    found = False
+    for student in base_cfg.multidistillation.students:
+        if rank in range(*student.ranks_range):
+            found = True
+            break
+    assert found, "rank of worker not in defined range"
+
+    name = student.name
+    config_path = student.config_path
+    n_gpus = student.ranks_range[1] - student.ranks_range[0]
+    assert global_batch_size % n_gpus == 0
+    total_n_gpus = distributed.get_world_size()
+
+    args.output_dir = os.path.join(base_output_dir, name)
+    args.opts += [f"train.output_dir={args.output_dir}"]
+    args.opts += [f"train.batch_size_per_gpu={global_batch_size // total_n_gpus}"]
+    args.config_file = os.path.abspath(config_path)
+    default_cfg = get_default_config()
+    cfg = OmegaConf.load(args.config_file)
+    cfg = OmegaConf.merge(default_cfg, cfg, base_cfg, OmegaConf.from_cli(args.opts))
+
+    global logger
+    setup_logging(output=args.output_dir, level=logging.INFO)
+
+    fix_random_seeds(seed + rank)
+
+    write_config(cfg, args.output_dir)
+    apply_scaling_rules_to_cfg(cfg)
+
+    return cfg
+
+
+def setup_job(
+    output_dir: Optional[str] = None,
+    distributed_enabled: bool = True,
+    logging_enabled: bool = True,
+    seed: Optional[int] = 0,
+    restrict_print_to_main_process: bool = True,
+    distributed_timeout: timedelta | None = None,
+):
+    """
+    Setup methods that should be done in every fairvit job
+    Initializes logging, distributed, random seeds and other utilities.
+    """
+    if output_dir is not None:
+        output_dir = os.path.realpath(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+
+    if logging_enabled:
+        setup_logging(
+            output=output_dir,
+            level=logging.INFO,
+            log_to_stdout_only_in_main_process=restrict_print_to_main_process,
+        )
+
+    if distributed_enabled:
+        distributed.enable(
+            overwrite=True,
+            nccl_async_error_handling=True,
+            restrict_print_to_main_process=restrict_print_to_main_process,
+            timeout=distributed_timeout,
+        )
+
+    if seed is not None:
+        rank = distributed.get_rank()
+        fix_random_seeds(seed + rank)
+
+    logger = logging.getLogger("dinov3")
+    logger.info("git:\n  {}\n".format(get_sha()))
+
+    # Log some python info
+    conda_env_name, conda_env_path = get_conda_env()
+    logger.info(f"conda env name: {conda_env_name}")
+    logger.info(f"conda env path: {conda_env_path}")
+    logger.info(f"python path: {sys.path}")
+
+
+def exit_job(distributed_enabled: bool = True, logging_enabled: bool = True):
+    if distributed_enabled:
+        distributed.disable()
+    if logging_enabled:
+        cleanup_logging()

dinov3/configs/ssl_default_config.yaml

@@ -0,0 +1,205 @@
+MODEL:
+  META_ARCHITECTURE: SSLMetaArch
+  DEVICE: cuda
+  WEIGHTS: ''
+  DTYPE: float32
+compute_precision:
+  param_dtype: bf16
+  reduce_dtype: fp32
+  sharding_strategy: SHARD_GRAD_OP
+dino:
+  loss_weight: 1.0
+  global_ignore_diagonal: true  # Whether to ignore A-A and B-B global pairs, default as in DINOv2, ignored by SSLMetaArch
+  head_n_prototypes: 65536
+  head_bottleneck_dim: 256
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 2048
+  koleo_loss_weight: 0.1
+  koleo_loss_distributed: false
+  koleo_topk: 1
+  koleo_distributed_replicas: 0
+  koleo_distributed_loss_group_size: null  # Size of the nearest neighbor set for distributed Koleo. If None, uses global batch size.
+  koleo_distributed_loss_group_data: true # group data from adjacent ranks to make sure koleo is applied on the same data distribution
+  force_weight_norm: false
+  reweight_dino_local_loss: false # If true, reweighting of DINO loss
+  local_loss_weight_schedule:  # Schedule for local loss weight, enabled if reweight_dino_local_loss is true
+    start: 0.5
+    peak: 0.5
+    end: 0.5
+    warmup_epochs: 0
+ibot:
+  loss_weight: 1.0
+  mask_sample_probability: 0.5
+  mask_ratio_min_max:
+  - 0.1
+  - 0.5
+  mask_random_circular_shift: false
+  force_masking_even_with_zero_weight: False
+  separate_head: true
+  head_n_prototypes: 65536
+  head_bottleneck_dim: 256
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 2048
+gram:
+  use_loss: false # (bool) if true gram is used, else not
+  compute_stats: false # (bool): if true compute auxilliary stats
+  loss_weight: 1.0 # (float): weight of the loss
+  ema_teacher: false # (bool): using the EMA teacher as GRAM teacher
+  ckpt: null  #(str): Checkpoint to the teacher
+  it_load_ema_teacher: -1 # (int): iteration at which the ema teacher is loaded into the gram teacher
+  rep_update: true # (bool): if true GRAM teacher updated every gram.update_frequency after iter gram.it_first_update steps
+  update_frequency: 50000 # (int): update frequency
+  it_first_update: 0 # (int): iteration of the first update
+  max_updates: null # (int): maximum number of updates to gram teacher. If None, it is unlimited
+  normalized: true # (bool): normalization of the features
+  img_level: false # (bool): if true GRAM computation at the image else, otherwise at the local batch level
+  remove_neg: false # (bool): if true remove the negative similarities before applying the loss
+  remove_only_teacher_neg: false # (bool): remove negative similarities of the teacher
+  tokens_used: all # (str): In [all, masked, unmasked]
+  global_teacher_resize_method: bicubic  # Method for resizing the outputs of the gram teacher
+  global_teacher_resize_antialias: false  # Whether to use antialiasing when resizing the outputs of the gram teacher
+  loss_weight_schedule: null  # (dict): If not None, use a schedule for the loss weight instead of `loss_weight`
+train:
+  batch_size_per_gpu: 64
+  dataset_path: ImageNet:split=TRAIN
+  data_config: null
+  output_dir: .
+  saveckp_freq: 20
+  seed: 0
+  num_workers: 10
+  OFFICIAL_EPOCH_LENGTH: 1250
+  monitor_gradient_norm: false
+  chunk_schedule: []
+  use_teacher_head: true
+  learn_from_teacher_tokens: false
+  centering: "sinkhorn_knopp" # or "sinkhorn_knopp"
+  checkpointing: false
+  checkpointing_full: false  # aggressive checkpointing
+  compile: true
+  cudagraphs: false
+  sharded_eval_checkpoint: false
+  cache_dataset: false
+student:
+  arch: vit_large
+  patch_size: 16
+  drop_path_rate: 0.3
+  layerscale: 1.0e-05
+  pretrained_weights: ''
+  ffn_layer: "mlp"
+  ffn_ratio: 4.0
+  resume_from_teacher_chkpt: ""
+  qkv_bias: true
+  proj_bias: true
+  ffn_bias: true
+  norm_layer: "layernorm"
+  n_storage_tokens: 0
+  mask_k_bias: false
+  untie_cls_and_patch_norms: false  # If true, use separate norms for CLS/reg and patch/mask tokens
+  untie_global_and_local_cls_norm: false  # If true, use separate norms for local and global crop CLS token during training
+  in_chans: 3
+  pos_embed_type: rope
+  pos_embed_rope_base: 100.0
+  pos_embed_rope_min_period: null
+  pos_embed_rope_max_period: null
+  pos_embed_rope_normalize_coords: separate  # min, max, separate
+  pos_embed_rope_shift_coords: null
+  pos_embed_rope_jitter_coords: null
+  pos_embed_rope_rescale_coords: null
+  pos_embed_rope_dtype: bf16
+  fp8_enabled: False  # Convert Linear layers to operate in fp8 precision
+  fp8_filter: "blocks"  # Regex that must appear in module path; empty means everything
+teacher:
+  momentum_teacher: 0.992
+  final_momentum_teacher: 1
+  warmup_teacher_temp: 0.04
+  teacher_temp: 0.07
+  warmup_teacher_temp_epochs: 30
+  in_chans: 3
+distillation:  # teacher
+  enabled: false
+  full_cfg_path: ""
+  checkpoint_path: ""
+multidistillation:
+  enabled: false
+hrft:  # non-hrft'd student
+  enabled: false
+  checkpoint_path: ""  # teacher_checkpoint path
+optim:
+  epochs: 100
+  optimizer: adamw
+  weight_decay: 0.04
+  weight_decay_end: 0.4
+  lr: 0.001
+  warmup_epochs: 10
+  min_lr: 1.0e-06
+  schedule_trunc_extra: 0.0  # Compute the schedule for (1 + schedule_trunc_extra) steps and truncate, .25 is a good choice
+  clip_grad: 3.0
+  freeze_last_layer_epochs: 1
+  scaling_rule: sqrt_wrt_1024
+  patch_embed_lr_mult: 0.2
+  dino_head_wd_multiplier: 1.0
+  layerwise_decay: 0.9
+  multi_tensor_optim: true
+  dump_fsdp_weights_path: ""
+  adamw_beta1: 0.9
+  adamw_beta2: 0.999
+crops:
+  global_crops_scale:
+  - 0.32
+  - 1.0
+  local_crops_number: 8
+  local_crops_scale:
+  - 0.05
+  - 0.32
+  global_crops_size: 224
+  local_crops_size: 96
+  global_local_crop_pairs_ratios: 1.0
+  gram_teacher_crops_size: null  # If not None, return crops for gram teacher
+  localcrops_subset_of_globalcrops: false
+  share_color_jitter: false
+  horizontal_flips: true
+  gram_teacher_no_distortions: false  # If True, no distortions are applied to gram teacher crops
+  rgb_mean:
+  - 0.485
+  - 0.456
+  - 0.406
+  rgb_std:
+  - 0.229
+  - 0.224
+  - 0.225
+evaluation:
+  eval_period_iterations: 12500
+  low_freq_every: 5
+  config_files:  # Must be in fairvit/eval/configs
+    high_freq: benchmark_high_frequency.yaml  # More often
+    low_freq: benchmark_low_frequency.yaml  # Less often
+checkpointing:
+  period: 3750
+  max_to_keep: 3
+  keep_every: 99999999999999999  # Save a checkpoint every N iterations, regardless of max_to_keep and period
+
+# Example of constant schedules with schedules v2
+# # schedules:
+# #   lr:
+# #     start: 0.0
+# #     peak: 1e-3
+# #     end: 1e-6
+# #     warmup_epochs: 10
+# #     freeze_last_layer_epochs: 1
+# #   weight_decay:
+# #     start: 0.04
+# #     peak: 0.04
+# #     end: 0.04
+# #     warmup_epochs: 0
+# #   momentum:
+# #     start: 0.992
+# #     peak: 0.992
+# #     end: 0.992
+# #     warmup_epochs: 0
+# #   teacher_temp:
+# #     start: 0.04
+# #     peak: 0.07
+# #     end: 0.07
+# #     warmup_epochs: 30

dinov3/configs/train/dinov3_vit7b16_gram_anchor.yaml

@@ -0,0 +1,203 @@
+MODEL:
+  META_ARCHITECTURE: SSLMetaArch
+  DEVICE: cuda
+  WEIGHTS: ''
+  DTYPE: float32
+compute_precision:
+  param_dtype: bf16
+  reduce_dtype: fp32
+  sharding_strategy: SHARD_GRAD_OP
+dino:
+  loss_weight: 1.0
+  global_ignore_diagonal: true
+  head_n_prototypes: 262144
+  head_bottleneck_dim: 512
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 8192
+  koleo_loss_weight: 0.1
+  koleo_loss_distributed: false
+  koleo_topk: 1
+  koleo_distributed_replicas: 0
+  koleo_distributed_loss_group_size: null
+  koleo_distributed_loss_group_data: true
+  force_weight_norm: false
+  reweight_dino_local_loss: true
+  local_loss_weight_schedule:
+    start: 1
+    peak: 1
+    end: 0.5
+    warmup_epochs: 1000
+    cosine_epochs: 1
+ibot:
+  loss_weight: 1.0
+  mask_sample_probability: 0.5
+  mask_ratio_min_max:
+  - 0.1
+  - 0.5
+  mask_random_circular_shift: false
+  force_masking_even_with_zero_weight: false
+  separate_head: true
+  head_n_prototypes: 98304
+  head_bottleneck_dim: 384
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 4096
+gram:
+  use_loss: true
+  compute_stats: false
+  loss_weight: 1.0
+  ema_teacher: false
+  ckpt: ignore
+  it_load_ema_teacher: -1
+  rep_update: true
+  update_frequency: 10000
+  it_first_update: 1010000
+  max_updates: 3
+  normalized: true
+  img_level: true
+  remove_neg: false
+  remove_only_teacher_neg: false
+  tokens_used: all
+  global_teacher_resize_method: bicubic
+  global_teacher_resize_antialias: false
+  loss_weight_schedule:
+    start: 0
+    peak: 0
+    end: 2.0
+    warmup_epochs: 1000
+    cosine_epochs: 1
+train:
+  batch_size_per_gpu: 16
+  dataset_path: null
+  saveckp_freq: 20
+  seed: 0
+  num_workers: 10
+  OFFICIAL_EPOCH_LENGTH: 1000
+  monitor_gradient_norm: false
+  chunk_schedule: []
+  cache_dataset: true
+  use_teacher_head: true
+  learn_from_teacher_tokens: false
+  centering: sinkhorn_knopp
+  checkpointing: true
+  checkpointing_full: true
+  compile: true
+  cudagraphs: false
+  cell_augmentation: false
+  cell_augmentation_type: hpa
+  sharded_eval_checkpoint: true
+student:
+  arch: vit_7b
+  patch_size: 16
+  drop_path_rate: 0.4
+  layerscale: 1.0e-05
+  patch_drop: 0.0
+  pretrained_weights: ''
+  ffn_layer: swiglu64
+  ffn_ratio: 3
+  resume_from_teacher_chkpt: ''
+  qkv_bias: false
+  proj_bias: true
+  ffn_bias: true
+  norm_layer: layernormbf16
+  n_storage_tokens: 4
+  untie_cls_and_patch_norms: false
+  untie_global_and_local_cls_norm: true
+  mask_k_bias: true
+  in_chans: 3
+  pos_embed_type: rope
+  pos_embed_rope_base: 100
+  pos_embed_rope_min_period: null
+  pos_embed_rope_max_period: null
+  pos_embed_rope_normalize_coords: separate
+  pos_embed_rope_shift_coords: null
+  pos_embed_rope_jitter_coords: null
+  pos_embed_rope_rescale_coords: 2
+  pos_embed_rope_dtype: fp32
+  fp8_enabled: true
+  fp8_filter: blocks
+teacher:
+  momentum_teacher: null
+  final_momentum_teacher: null
+  warmup_teacher_temp: null
+  teacher_temp: null
+  warmup_teacher_temp_epochs: null
+  in_chans: 3
+distillation:
+  enabled: false
+  full_cfg_path: ''
+  checkpoint_path: ''
+multidistillation:
+  enabled: false
+hrft:
+  enabled: false
+  checkpoint_path: ''
+optim:
+  epochs: 1200
+  optimizer: adamw
+  weight_decay: null
+  weight_decay_end: null
+  lr: null
+  warmup_epochs: null
+  min_lr: null
+  schedule_trunc_extra: null
+  clip_grad: 30.0
+  freeze_last_layer_epochs: null
+  scaling_rule: sqrt_wrt_1024
+  patch_embed_lr_mult: 0.2
+  dino_head_wd_multiplier: 1.0
+  layerwise_decay: 0.98
+  multi_tensor_optim: true
+  dump_fsdp_weights_path: ''
+  adamw_beta1: 0.9
+  adamw_beta2: 0.99
+crops:
+  global_crops_scale:
+  - 0.32
+  - 1.0
+  local_crops_number: 8
+  local_crops_scale:
+  - 0.05
+  - 0.32
+  global_crops_size: 256
+  local_crops_size: 112
+  gram_teacher_crops_size: 512
+  localcrops_subset_of_globalcrops: false
+  share_color_jitter: false
+  horizontal_flips: false
+  gram_teacher_no_distortions: true
+  rgb_mean:
+  - 0.485
+  - 0.456
+  - 0.406
+  rgb_std:
+  - 0.229
+  - 0.224
+  - 0.225
+checkpointing:
+  period: 1000
+  max_to_keep: 3
+  keep_every: 50000
+schedules:
+  lr:
+    start: 0
+    peak: 3.0e-05
+    end: 3.0e-05
+    warmup_epochs: 100
+    freeze_last_layer_epochs: 5
+  weight_decay:
+    start: 0.04
+    peak: 0.04
+    end: 0.04
+    warmup_epochs: 0
+  teacher_temp:
+    start: 0.04
+    peak: 0.07
+    end: 0.07
+    warmup_epochs: 100
+  momentum:
+    start: 0.999
+    peak: 0.999
+    end: 0.999
+    warmup_epochs: 0

dinov3/configs/train/dinov3_vit7b16_high_res_adapt.yaml

@@ -0,0 +1,224 @@
+MODEL:
+  META_ARCHITECTURE: SSLMetaArch
+  DEVICE: cuda
+  WEIGHTS: ''
+  DTYPE: float32
+compute_precision:
+  param_dtype: bf16
+  reduce_dtype: fp32
+  sharding_strategy: SHARD_GRAD_OP
+dino:
+  loss_weight: 1.0
+  global_ignore_diagonal: true
+  head_n_prototypes: 262144
+  head_bottleneck_dim: 512
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 8192
+  koleo_loss_weight: 0.1
+  koleo_loss_distributed: true
+  koleo_topk: 1
+  koleo_distributed_replicas: 0
+  koleo_distributed_loss_group_size: 16
+  force_weight_norm: false
+  reweight_dino_local_loss: true
+  local_loss_weight_schedule:
+    start: 0.5
+    peak: 0.5
+    end: 0.5
+    warmup_epochs: 0
+    cosine_epochs: 0
+  koleo_distributed_loss_group_data: true
+ibot:
+  loss_weight: 1.0
+  mask_sample_probability: 0.5
+  mask_ratio_min_max:
+  - 0.1
+  - 0.5
+  mask_random_circular_shift: false
+  force_masking_even_with_zero_weight: false
+  separate_head: true
+  head_n_prototypes: 98304
+  head_bottleneck_dim: 384
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 4096
+gram:
+  use_loss: true
+  compute_stats: false
+  loss_weight: 1.0
+  ema_teacher: false
+  it_load_ema_teacher: -1
+  rep_update: false
+  update_frequency: 10000
+  it_first_update: 1010000
+  max_updates: 3
+  normalized: true
+  img_level: true
+  remove_neg: false
+  remove_only_teacher_neg: false
+  tokens_used: all
+  global_teacher_resize_method: bicubic
+  global_teacher_resize_antialias: false
+  loss_weight_schedule:
+    start: 1.5
+    peak: 1.5
+    end: 1.5
+    warmup_epochs: 0
+    cosine_epochs: 0
+train:
+  batch_size_per_gpu: 8
+  dataset_path: null
+  saveckp_freq: 20
+  seed: 0
+  num_workers: 2
+  OFFICIAL_EPOCH_LENGTH: 1000
+  monitor_gradient_norm: false
+  chunk_schedule: []
+  cache_dataset: true
+  use_teacher_head: true
+  learn_from_teacher_tokens: false
+  centering: sinkhorn_knopp
+  checkpointing: true
+  checkpointing_full: true
+  compile: true
+  cudagraphs: false
+  cell_augmentation: false
+  cell_augmentation_type: hpa
+  sharded_eval_checkpoint: true
+student:
+  arch: vit_7b
+  patch_size: 16
+  drop_path_rate: 0.4
+  layerscale: 1.0e-05
+  patch_drop: 0.0
+  pretrained_weights: ''
+  ffn_layer: swiglu64
+  ffn_ratio: 3
+  resume_from_teacher_chkpt: ''
+  qkv_bias: false
+  proj_bias: true
+  ffn_bias: true
+  norm_layer: layernormbf16
+  n_storage_tokens: 4
+  untie_cls_and_patch_norms: false
+  untie_global_and_local_cls_norm: true
+  mask_k_bias: true
+  in_chans: 3
+  pos_embed_type: rope
+  pos_embed_rope_base: 100
+  pos_embed_rope_min_period: null
+  pos_embed_rope_max_period: null
+  pos_embed_rope_normalize_coords: separate
+  pos_embed_rope_shift_coords: null
+  pos_embed_rope_jitter_coords: null
+  pos_embed_rope_rescale_coords: 2
+  pos_embed_rope_dtype: fp32
+  fp8_enabled: true
+  fp8_filter: blocks
+teacher:
+  momentum_teacher: null
+  final_momentum_teacher: null
+  warmup_teacher_temp: null
+  teacher_temp: null
+  warmup_teacher_temp_epochs: null
+  in_chans: 3
+distillation:
+  enabled: false
+  full_cfg_path: ''
+  checkpoint_path: ''
+multidistillation:
+  enabled: false
+hrft:
+  enabled: false
+  checkpoint_path: ''
+optim:
+  epochs: 30
+  optimizer: adamw
+  weight_decay: null
+  weight_decay_end: null
+  lr: null
+  warmup_epochs: null
+  min_lr: null
+  schedule_trunc_extra: null
+  clip_grad: 30.0
+  freeze_last_layer_epochs: null
+  scaling_rule: sqrt_wrt_1024
+  patch_embed_lr_mult: 0.2
+  dino_head_wd_multiplier: 1.0
+  layerwise_decay: 0.98
+  multi_tensor_optim: true
+  dump_fsdp_weights_path: ''
+  adamw_beta1: 0.9
+  adamw_beta2: 0.99
+crops:
+  global_crops_scale:
+  - 0.32
+  - 1.0
+  local_crops_number: 8
+  local_crops_scale:
+  - 0.05
+  - 0.32
+  global_crops_size:
+  - 512
+  - 768
+  - 768
+  - 768
+  - 768
+  local_crops_size:
+  - 112
+  - 112
+  - 168
+  - 224
+  - 336
+  global_local_crop_pairs_ratios:
+  - 0.3
+  - 0.3
+  - 0.3
+  - 0.05
+  - 0.05
+  gram_teacher_crops_size:
+  - 768
+  - 1152
+  - 1152
+  - 1152
+  - 1152
+  localcrops_subset_of_globalcrops: false
+  share_color_jitter: false
+  horizontal_flips: false
+  gram_teacher_no_distortions: true
+  rgb_mean:
+  - 0.485
+  - 0.456
+  - 0.406
+  rgb_std:
+  - 0.229
+  - 0.224
+  - 0.225
+checkpointing:
+  period: 250
+  max_to_keep: 3
+  keep_every: 50000
+schedules:
+  lr:
+    start: 0
+    peak: 0
+    end: 1.25e-05
+    warmup_epochs: 0
+    freeze_last_layer_epochs: 0
+    cosine_epochs: 10
+  weight_decay:
+    start: 0.04
+    peak: 0.04
+    end: 0.04
+    warmup_epochs: 0
+  teacher_temp:
+    start: 0.07
+    peak: 0.07
+    end: 0.07
+    warmup_epochs: 0
+  momentum:
+    start: 0.999
+    peak: 0.999
+    end: 0.999
+    warmup_epochs: 0

dinov3/configs/train/dinov3_vit7b16_pretrain.yaml

@@ -0,0 +1,172 @@
+MODEL:
+  META_ARCHITECTURE: SSLMetaArch
+  DEVICE: cuda
+  WEIGHTS: ''
+  DTYPE: float32
+compute_precision:
+  param_dtype: bf16
+  reduce_dtype: fp32
+  sharding_strategy: SHARD_GRAD_OP
+dino:
+  loss_weight: 1.0
+  global_ignore_diagonal: true
+  head_n_prototypes: 262144
+  head_bottleneck_dim: 512
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 8192
+  koleo_loss_weight: 0.1
+  koleo_loss_distributed: false
+  koleo_topk: 1
+  koleo_distributed_replicas: 0
+  koleo_distributed_loss_group_size: null
+  force_weight_norm: false
+ibot:
+  loss_weight: 1.0
+  mask_sample_probability: 0.5
+  mask_ratio_min_max:
+  - 0.1
+  - 0.5
+  mask_random_circular_shift: false
+  force_masking_even_with_zero_weight: false
+  separate_head: true
+  head_n_prototypes: 98304
+  head_bottleneck_dim: 384
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 4096
+gram:
+  use_loss: false
+  compute_stats: false
+train:
+  batch_size_per_gpu: 16
+  dataset_path: null
+  saveckp_freq: 20
+  seed: 0
+  num_workers: 10
+  OFFICIAL_EPOCH_LENGTH: 1000
+  monitor_gradient_norm: false
+  chunk_schedule: []
+  cache_dataset: true
+  use_teacher_head: true
+  learn_from_teacher_tokens: false
+  centering: sinkhorn_knopp
+  checkpointing: true
+  checkpointing_full: false
+  compile: true
+  cudagraphs: false
+  cell_augmentation: false
+  cell_augmentation_type: hpa
+  sharded_eval_checkpoint: true
+student:
+  arch: vit_7b
+  patch_size: 16
+  drop_path_rate: 0.4
+  layerscale: 1.0e-05
+  patch_drop: 0.0
+  pretrained_weights: ''
+  ffn_layer: swiglu64
+  ffn_ratio: 3
+  resume_from_teacher_chkpt: ''
+  qkv_bias: false
+  proj_bias: true
+  ffn_bias: true
+  norm_layer: layernormbf16
+  n_storage_tokens: 4
+  untie_cls_and_patch_norms: false
+  untie_global_and_local_cls_norm: true
+  mask_k_bias: true
+  in_chans: 3
+  pos_embed_type: rope
+  pos_embed_rope_base: 100
+  pos_embed_rope_min_period: null
+  pos_embed_rope_max_period: null
+  pos_embed_rope_normalize_coords: separate
+  pos_embed_rope_shift_coords: null
+  pos_embed_rope_jitter_coords: null
+  pos_embed_rope_rescale_coords: 2
+  pos_embed_rope_dtype: fp32
+  fp8_enabled: true
+  fp8_filter: blocks
+teacher:
+  momentum_teacher: null
+  final_momentum_teacher: null
+  warmup_teacher_temp: null
+  teacher_temp: null
+  warmup_teacher_temp_epochs: null
+  in_chans: 3
+distillation:
+  enabled: false
+  full_cfg_path: ''
+  checkpoint_path: ''
+multidistillation:
+  enabled: false
+hrft:
+  enabled: false
+  checkpoint_path: ''
+optim:
+  epochs: 1000
+  optimizer: adamw
+  weight_decay: null
+  weight_decay_end: null
+  lr: null
+  warmup_epochs: null
+  min_lr: null
+  schedule_trunc_extra: null
+  clip_grad: 30.0
+  freeze_last_layer_epochs: null
+  scaling_rule: sqrt_wrt_1024
+  patch_embed_lr_mult: 0.2
+  dino_head_wd_multiplier: 1.0
+  layerwise_decay: 0.98
+  multi_tensor_optim: true
+  dump_fsdp_weights_path: ''
+  adamw_beta1: 0.9
+  adamw_beta2: 0.99
+crops:
+  global_crops_scale:
+  - 0.32
+  - 1.0
+  local_crops_number: 8
+  local_crops_scale:
+  - 0.05
+  - 0.32
+  global_crops_size: 256
+  local_crops_size: 112
+  localcrops_subset_of_globalcrops: false
+  share_color_jitter: false
+  horizontal_flips: false
+  rgb_mean:
+  - 0.485
+  - 0.456
+  - 0.406
+  rgb_std:
+  - 0.229
+  - 0.224
+  - 0.225
+checkpointing:
+  period: 1000
+  max_to_keep: 3
+  keep_every: 50000
+schedules:
+  lr:
+    start: 0
+    peak: 5.0e-05
+    end: 5.0e-05
+    warmup_epochs: 100
+    freeze_last_layer_epochs: 5
+  weight_decay:
+    start: 0.04
+    peak: 0.04
+    end: 0.04
+    warmup_epochs: 0
+  teacher_temp:
+    start: 0.04
+    peak: 0.07
+    end: 0.07
+    warmup_epochs: 100
+  momentum:
+    start: 0.994
+    peak: 0.994
+    end: 0.994
+    warmup_epochs: 0

dinov3/configs/train/dinov3_vitl16_lvd1689m_distilled.yaml

@@ -0,0 +1,251 @@
+MODEL:
+  META_ARCHITECTURE: MultiDistillationMetaArch
+  DEVICE: cuda
+  WEIGHTS: ''
+  DTYPE: float32
+compute_precision:
+  param_dtype: bf16
+  reduce_dtype: fp32
+  sharding_strategy: SHARD_GRAD_OP
+dino:
+  loss_weight: 1.0
+  global_ignore_diagonal: true
+  head_n_prototypes: 262144
+  head_bottleneck_dim: 512
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 8192
+  koleo_loss_weight: 0.1
+  koleo_loss_distributed: false
+  koleo_topk: 1
+  koleo_distributed_replicas: 0
+  koleo_distributed_loss_group_size: null
+  koleo_distributed_loss_group_data: true
+  force_weight_norm: false
+  reweight_dino_local_loss: false
+  local_loss_weight_schedule:
+    start: 0.5
+    peak: 0.5
+    end: 0.5
+    warmup_epochs: 0
+ibot:
+  loss_weight: 1.0
+  mask_sample_probability: 0.5
+  mask_ratio_min_max:
+  - 0.1
+  - 0.5
+  mask_random_circular_shift: false
+  force_masking_even_with_zero_weight: false
+  separate_head: true
+  head_n_prototypes: 98304
+  head_bottleneck_dim: 384
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 4096
+coding_rate_loss:
+  use_cls_loss: false
+  cls_loss_weight: 0.2
+  use_masked_patches_loss: false
+  masked_patches_loss_weight: 0.1
+  epsilon: 8
+gram:
+  use_loss: false
+  compute_stats: false
+  loss_weight: 1.0
+  ema_teacher: false
+  ckpt: null
+  it_load_ema_teacher: -1
+  rep_update: true
+  update_frequency: 50000
+  it_first_update: 0
+  max_updates: null
+  normalized: true
+  img_level: false
+  remove_neg: false
+  remove_only_teacher_neg: false
+  tokens_used: all
+  global_teacher_resize_method: bicubic
+  global_teacher_resize_antialias: false
+  loss_weight_schedule: null
+train:
+  batch_size_per_gpu: 3
+  dataset_path: <TRAIN/DATASET>
+  output_dir: <OUTPUT/DIR>
+  saveckp_freq: 20
+  seed: 0
+  num_workers: 2
+  OFFICIAL_EPOCH_LENGTH: 1250
+  monitor_gradient_norm: false
+  chunk_schedule: []
+  cache_dataset: true
+  use_teacher_head: true
+  learn_from_teacher_tokens: false
+  centering: sinkhorn_knopp
+  checkpointing: true
+  checkpointing_full: true
+  compile: true
+  cudagraphs: false
+  cell_augmentation: false
+  cell_augmentation_type: hpa
+  sharded_eval_checkpoint: false
+student:
+  arch: vit_large
+  patch_size: 16
+  drop_path_rate: 0.0
+  layerscale: 1.0e-05
+  drop_path_uniform: true
+  drop_path_shape: uniform
+  patch_drop: 0.0
+  pretrained_weights: ''
+  sin_cos_embeddings: false
+  fourier_embeddings: false
+  fourier_encoding_dim: 64
+  multiple_pos_embeddings: false
+  cls_pos_embedding: false
+  reg_pos_embedding: false
+  ffn_layer: mlp
+  ffn_ratio: 4.0
+  resume_from_teacher_chkpt: <PATH/TO/HRFT/TEACHER>
+  block_chunks: 0
+  qkv_bias: true
+  proj_bias: true
+  ffn_bias: true
+  norm_layer: layernormbf16
+  n_storage_tokens: 4
+  mask_attention: false
+  mask_register_attention: false
+  untie_cls_and_patch_norms: false
+  untie_global_and_local_cls_norm: false
+  interpolate_offset: 0.0
+  interpolate_antialias: true
+  mask_k_bias: true
+  init_std_cls: 0.02
+  init_std_reg: 0.02
+  rescale_weights_by_layer_id: false
+  in_chans: 3
+  pos_embed_grid_size: 48
+  pos_embed_type: ropenew
+  pos_embed_rope_gamma: 1.0
+  pos_embed_rope_init_multi_frequencies: false
+  pos_embed_rope_base: 100
+  pos_embed_rope_min_period: null
+  pos_embed_rope_max_period: null
+  pos_embed_rope_normalize_coords: separate
+  pos_embed_rope_shift_coords: null
+  pos_embed_rope_jitter_coords: null
+  pos_embed_rope_rescale_coords: 2
+  pos_embed_rope_dtype: bf16
+  sparse24_ranges: []
+  sparse24_filter:
+  - mlp
+  sparse24_default: false
+  fp8_enabled: false
+  fp8_filter: blocks
+teacher:
+  momentum_teacher: 0.994
+  final_momentum_teacher: 1
+  warmup_teacher_temp: 0.04
+  teacher_temp: 0.07
+  warmup_teacher_temp_epochs: 120
+  in_chans: 3
+distillation:
+  enabled: true
+  full_cfg_path: <PATH/TO/TEACHER/CONFIG/config.yaml>
+  checkpoint_path: <PATH/TO/TEACHER/checkpoint.pth>
+multidistillation:
+  enabled: true
+  global_batch_size: 1920
+  students:
+  - name: vits_mlp4_4
+    config_path:  <PATH/TO/STUDENT/CONFIG/vits_mlp4_4.yaml>
+    ranks_range:
+    - 0
+    - 48
+  - name: vitsp_swiglu6_1
+    config_path: <PATH/TO/STUDENT/CONFIG/vitsp_swiglu6_1.yaml>
+    ranks_range:
+    - 48
+    - 96
+  - name: vitb_mlp4_3
+    config_path: <PATH/TO/STUDENT/CONFIG/vitb_mlp4_3.yaml>
+    ranks_range:
+    - 96
+    - 176
+  - name: vitl_mlp4_1
+    config_path:  <PATH/TO/STUDENT/CONFIG/vitl_mlp4_1.yaml>
+    ranks_range:
+    - 176
+    - 296
+hrft:
+  enabled: false
+  checkpoint_path: ''
+optim:
+  epochs: 20
+  optimizer: adamw
+  weight_decay: 0.04
+  weight_decay_end: 0.2
+  lr: 0.0002
+  warmup_epochs: 0
+  min_lr: 1.0e-06
+  schedule_trunc_extra: 0.0
+  clip_grad: 3.0
+  freeze_last_layer_epochs: 0
+  scaling_rule: sqrt_wrt_1024
+  patch_embed_lr_mult: 0.2
+  dino_head_wd_multiplier: 1.0
+  layerwise_decay: 0.99
+  multi_tensor_optim: true
+  dump_fsdp_weights_path: ''
+  adamw_beta1: 0.9
+  adamw_beta2: 0.999
+crops:
+  global_crops_scale:
+  - 0.32
+  - 1.0
+  local_crops_number: 8
+  local_crops_scale:
+  - 0.05
+  - 0.32
+  global_crops_size: 256
+  local_crops_size: 112
+  global_local_crop_pairs_ratios: 1.0
+  gram_teacher_crops_size: 256
+  localcrops_subset_of_globalcrops: false
+  share_color_jitter: false
+  horizontal_flips: false
+  gram_teacher_no_distortions: false
+  rgb_mean:
+  - 0.485
+  - 0.456
+  - 0.406
+  rgb_std:
+  - 0.229
+  - 0.224
+  - 0.225
+checkpointing:
+  period: 3750
+  max_to_keep: 3
+  keep_every: 99999999999999999
+schedules:
+  weight_decay:
+    start: 0.04
+    peak: 0.04
+    end: 0.04
+    warmup_epochs: 0
+  teacher_temp:
+    start: 0.04
+    peak: 0.07
+    end: 0.07
+    warmup_epochs: 0
+  lr:
+    start: 0
+    peak: 0
+    end: 5.0e-05
+    warmup_epochs: 0
+    freeze_last_layer_epochs: 0
+    cosine_epochs: 10
+  momentum:
+    start: 0.994
+    peak: 0.994
+    end: 1.0
+    warmup_epochs: 0

dinov3/configs/train/multi_distillation_test.yaml

@@ -0,0 +1,27 @@
+MODEL:
+  META_ARCHITECTURE: MultiDistillationMetaArch
+multidistillation:
+  enabled: true
+  global_batch_size: 256
+  students:
+  - name: vits
+    config_path: dinov3/configs/train/multidist_tests/vits_p16.yaml
+    ranks_range:
+    - 0
+    - 4
+  - name: vitb
+    config_path: dinov3/configs/train/multidist_tests/vitb_p16.yaml
+    ranks_range:
+    - 4
+    - 8
+distillation:  # teacher
+  enabled: true
+  full_cfg_path: dinov3/configs/train/vitl_im1k_lin834.yaml
+  checkpoint_path: ignore
+train:
+  dataset_path: ImageNet:split=TRAIN
+  cache_dataset: false
+  centering: "sinkhorn_knopp"
+  compile: true
+ibot:
+  separate_head: true

dinov3/configs/train/multidist_tests/vitb_p16.yaml

@@ -0,0 +1,7 @@
+# this corresponds to the default config
+train:
+  dataset_path: ImageNet:split=TRAIN
+  checkpointing: true
+student:
+  drop_path_rate: 0.1
+  arch: vit_base

dinov3/configs/train/multidist_tests/vits_p16.yaml

@@ -0,0 +1,6 @@
+# this corresponds to the default config
+train:
+  dataset_path: ImageNet:split=TRAIN
+student:
+  drop_path_rate: 0.1
+  arch: vit_small

dinov3/configs/train/vitl_im1k_lin834.yaml

@@ -0,0 +1,143 @@
+# tested on RSC: /checkpoint/dino/qas/rope/vitl16_im1k/
+# gives 82.2 im1k-knn, 83.3 im1k-linear
+# runs with a total batch size of 2048 (64/gpu, 4 nodes here)
+# runs at 0.57s/iter
+MODEL:
+  META_ARCHITECTURE: SSLMetaArch
+  DEVICE: cuda
+  WEIGHTS: ''
+  DTYPE: float32
+compute_precision:
+  param_dtype: bf16
+  reduce_dtype: fp32
+  sharding_strategy: SHARD_GRAD_OP
+dino:
+  loss_weight: 1.0
+  global_ignore_diagonal: true
+  head_n_prototypes: 65536
+  head_bottleneck_dim: 256
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 2048
+  koleo_loss_weight: 0.1
+  koleo_loss_distributed: false
+  koleo_topk: 1
+  koleo_distributed_replicas: 0
+  force_weight_norm: false
+ibot:
+  loss_weight: 1.0
+  mask_sample_probability: 0.5
+  mask_ratio_min_max:
+  - 0.1
+  - 0.5
+  mask_random_circular_shift: false
+  force_masking_even_with_zero_weight: false
+  separate_head: true
+  head_n_prototypes: 65536
+  head_bottleneck_dim: 256
+  head_norm_last_layer: false
+  head_nlayers: 3
+  head_hidden_dim: 2048
+train:
+  batch_size_per_gpu: 64
+  dataset_path: ImageNet:split=TRAIN
+  output_dir: /checkpoint/dino/qas/rope/vitl16_im1k
+  saveckp_freq: 20
+  seed: 0
+  num_workers: 10
+  OFFICIAL_EPOCH_LENGTH: 1250
+  monitor_gradient_norm: false
+  chunk_schedule: []
+  cache_dataset: true
+  use_teacher_head: true
+  learn_from_teacher_tokens: false
+  centering: sinkhorn_knopp
+  checkpointing: false
+  compile: true
+  cudagraphs: false
+  cell_augmentation: false
+  cell_augmentation_type: hpa
+student:
+  arch: vit_large
+  patch_size: 16
+  drop_path_rate: 0.3
+  layerscale: 1.0e-05
+  patch_drop: 0.0
+  pretrained_weights: ''
+  ffn_layer: mlp
+  ffn_ratio: 4.0
+  resume_from_teacher_chkpt: ''
+  qkv_bias: true
+  proj_bias: true
+  ffn_bias: true
+  norm_layer: layernorm
+  n_storage_tokens: 0
+  mask_k_bias: false
+  in_chans: 3
+  pos_embed_type: rope
+  pos_embed_rope_base: 100.0
+  pos_embed_rope_min_period: null
+  pos_embed_rope_max_period: null
+  pos_embed_rope_normalize_coords: separate  # min, max, separate
+  pos_embed_rope_shift_coords: null
+  pos_embed_rope_jitter_coords: null
+  pos_embed_rope_rescale_coords: null
+  pos_embed_rope_dtype: bf16
+  fp8_enabled: False  # Convert Linear layers to operate in fp8 precision
+  fp8_filter: "blocks"  # Regex that must appear in module path; empty means everything
+teacher:
+  momentum_teacher: 0.992
+  final_momentum_teacher: 1
+  warmup_teacher_temp: 0.04
+  teacher_temp: 0.07
+  warmup_teacher_temp_epochs: 30
+  in_chans: 3
+distillation:
+  enabled: false
+  full_cfg_path: ''
+  checkpoint_path: ''
+multidistillation:
+  enabled: false
+hrft:
+  enabled: false
+  checkpoint_path: ''
+optim:
+  epochs: 100
+  optimizer: adamw
+  weight_decay: 0.04
+  weight_decay_end: 0.4
+  lr: 0.001
+  warmup_epochs: 10
+  min_lr: 1.0e-06
+  clip_grad: 3.0
+  freeze_last_layer_epochs: 1
+  scaling_rule: sqrt_wrt_1024
+  patch_embed_lr_mult: 0.2
+  dino_head_wd_multiplier: 1.0
+  layerwise_decay: 0.9
+  multi_tensor_optim: true
+  dump_fsdp_weights_path: ''
+  adamw_beta1: 0.9
+  adamw_beta2: 0.999
+crops:
+  global_crops_scale:
+  - 0.32
+  - 1.0
+  local_crops_number: 8
+  local_crops_scale:
+  - 0.05
+  - 0.32
+  global_crops_size: 224
+  local_crops_size: 96
+  localcrops_subset_of_globalcrops: false
+  share_color_jitter: false
+  horizontal_flips: true
+evaluation:
+  eval_period_iterations: 12500
+  low_freq_every: 5
+  config_files:
+    high_freq: benchmark_high_frequency.yaml
+    low_freq: benchmark_low_frequency.yaml
+checkpointing:
+  period: 3750
+  max_to_keep: 3

dinov3/data/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from .adapters import DatasetWithEnumeratedTargets
+from .augmentations import DataAugmentationDINO
+from .collate import collate_data_and_cast
+from .loaders import SamplerType, make_data_loader, make_dataset
+from .meta_loaders import CombinedDataLoader
+from .masking import MaskingGenerator
+from .transforms import make_classification_eval_transform, make_classification_train_transform

dinov3/data/adapters.py

@@ -0,0 +1,68 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from typing import Any, Optional, Tuple
+
+from torch.utils.data import Dataset
+
+
+def extend_samples_with_index(dataset_class):
+    class DatasetWithIndex(dataset_class):
+        def __init__(self, **kwargs) -> None:
+            root = dataset_class.get_root()
+            super().__init__(root=root, **kwargs)
+
+        def __getitem__(self, index: int):
+            image, target = super().__getitem__(index)
+            return image, target, index
+
+    return DatasetWithIndex
+
+
+class DatasetWithEnumeratedTargets(Dataset):
+    """
+    If pad_dataset is set, pads based on torch's DistributedSampler implementation, which
+    with drop_last=False pads the last batch to be a multiple of the world size.
+    https://github.com/pytorch/pytorch/blob/main/torch/utils/data/distributed.py#L91
+    """
+
+    def __init__(self, dataset: Dataset, pad_dataset: bool = False, num_replicas: Optional[int] = None):
+        self._dataset = dataset
+        self._size = len(self._dataset)
+        self._padded_size = self._size
+        self._pad_dataset = pad_dataset
+        if self._pad_dataset:
+            assert num_replicas is not None, "num_replicas should be set if pad_dataset is True"
+            self._padded_size = num_replicas * ((len(dataset) + num_replicas - 1) // num_replicas)
+
+    def get_image_relpath(self, index: int) -> str:
+        assert self._pad_dataset or index < self._size
+        return self._dataset.get_image_relpath(index % self._size)
+
+    def get_image_data(self, index: int) -> bytes:
+        assert self._pad_dataset or index < self._size
+        return self._dataset.get_image_data(index % self._size)
+
+    def get_target(self, index: int) -> Tuple[Any, int]:
+        target = self._dataset.get_target(index % self._size)
+        if index >= self._size:
+            assert self._pad_dataset
+            return (-1, target)
+        return (index, target)
+
+    def get_sample_decoder(self, index: int) -> Any:
+        assert self._pad_dataset or index < self._size
+        return self._dataset.get_sample_decoder(index % self._size)
+
+    def __getitem__(self, index: int) -> Tuple[Any, Tuple[Any, int]]:
+        image, target = self._dataset[index % self._size]
+        if index >= self._size:
+            assert self._pad_dataset
+            return image, (-1, target)
+        target = index if target is None else target
+        return image, (index, target)
+
+    def __len__(self) -> int:
+        return self._padded_size

dinov3/data/augmentations.py

@@ -0,0 +1,227 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import logging
+
+import numpy as np
+import torch
+from torch import nn
+from torchvision.transforms import v2
+
+from dinov3.data.transforms import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, GaussianBlur, make_normalize_transform
+
+logger = logging.getLogger("dinov3")
+
+
+class DataAugmentationDINO(object):
+    def __init__(
+        self,
+        global_crops_scale,
+        local_crops_scale,
+        local_crops_number,
+        global_crops_size=224,
+        local_crops_size=96,
+        gram_teacher_crops_size=None,
+        gram_teacher_no_distortions=False,
+        teacher_no_color_jitter=False,
+        local_crops_subset_of_global_crops=False,
+        patch_size=16,
+        share_color_jitter=False,
+        horizontal_flips=True,
+        mean=IMAGENET_DEFAULT_MEAN,
+        std=IMAGENET_DEFAULT_STD,
+    ):
+        self.global_crops_scale = global_crops_scale
+        self.local_crops_scale = local_crops_scale
+        self.local_crops_number = local_crops_number
+        self.global_crops_size = global_crops_size
+        self.local_crops_size = local_crops_size
+        self.gram_teacher_crops_size = gram_teacher_crops_size
+        self.gram_teacher_no_distortions = gram_teacher_no_distortions
+        self.teacher_no_color_jitter = teacher_no_color_jitter
+        self.local_crops_subset_of_global_crops = local_crops_subset_of_global_crops
+        self.patch_size = patch_size
+        self.share_color_jitter = share_color_jitter
+        self.mean = mean
+        self.std = std
+
+        logger.info("###################################")
+        logger.info("Using data augmentation parameters:")
+        logger.info(f"global_crops_scale: {global_crops_scale}")
+        logger.info(f"local_crops_scale: {local_crops_scale}")
+        logger.info(f"local_crops_number: {local_crops_number}")
+        logger.info(f"global_crops_size: {global_crops_size}")
+        logger.info(f"local_crops_size: {local_crops_size}")
+        logger.info(f"gram_crops_size: {gram_teacher_crops_size}")
+        logger.info(f"gram_teacher_no_distortions: {gram_teacher_no_distortions}")
+        logger.info(f"teacher_no_color_jitter: {teacher_no_color_jitter}")
+        logger.info(f"local_crops_subset_of_global_crops: {local_crops_subset_of_global_crops}")
+        logger.info(f"patch_size if local_crops_subset_of_global_crops: {patch_size}")
+        logger.info(f"share_color_jitter: {share_color_jitter}")
+        logger.info(f"horizontal flips: {horizontal_flips}")
+        logger.info("###################################")
+
+        # Global crops and gram teacher crops can have different sizes. We first take a crop of the maximum size
+        # and then resize it to the desired size for global and gram teacher crops.
+        global_crop_max_size = max(global_crops_size, gram_teacher_crops_size if gram_teacher_crops_size else 0)
+
+        # random resized crop and flip
+        self.geometric_augmentation_global = v2.Compose(
+            [
+                v2.RandomResizedCrop(
+                    global_crop_max_size,
+                    scale=global_crops_scale,
+                    interpolation=v2.InterpolationMode.BICUBIC,
+                ),
+                v2.RandomHorizontalFlip(p=0.5 if horizontal_flips else 0.0),
+            ]
+        )
+
+        resize_global = nn.Identity()  # Resize transform applied to global crops after random crop
+        self.resize_global_post_transf = (
+            nn.Identity()
+        )  # Resize transform applied to global crops after all other transforms
+        self.resize_gram_teacher = None  # Resize transform applied to crops for gram teacher
+        if gram_teacher_crops_size is not None:
+            # All resize transforms will do nothing if the crop size is already the desired size.
+            if gram_teacher_no_distortions:
+                # When there a no distortions for the gram teacher crop, we can resize before the distortions.
+                # This is the preferred order, because it keeps the image size for the augmentations consistent,
+                # which matters e.g. for GaussianBlur.
+                resize_global = v2.Resize(
+                    global_crops_size,
+                    interpolation=v2.InterpolationMode.BICUBIC,
+                )
+            else:
+                # When there a no distortions for the gram teacher crop, we need to resize after the distortions,
+                # because the distortions are shared between global and gram teacher crops.
+                self.resize_global_post_transf = v2.Resize(
+                    global_crops_size,
+                    interpolation=v2.InterpolationMode.BICUBIC,
+                )
+
+            self.resize_gram_teacher = v2.Resize(
+                gram_teacher_crops_size,
+                interpolation=v2.InterpolationMode.BICUBIC,
+            )
+
+        self.geometric_augmentation_local = v2.Compose(
+            [
+                v2.RandomResizedCrop(
+                    local_crops_size,
+                    scale=local_crops_scale,
+                    interpolation=v2.InterpolationMode.BICUBIC,
+                ),
+                v2.RandomHorizontalFlip(p=0.5 if horizontal_flips else 0.0),
+            ]
+        )
+
+        # color distortions / blurring
+        color_jittering = v2.Compose(
+            [
+                v2.RandomApply(
+                    [v2.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1)],
+                    p=0.8,
+                ),
+                v2.RandomGrayscale(p=0.2),
+            ]
+        )
+
+        global_transfo1_extra = GaussianBlur(p=1.0)
+
+        global_transfo2_extra = v2.Compose(
+            [
+                GaussianBlur(p=0.1),
+                v2.RandomSolarize(threshold=128, p=0.2),
+            ]
+        )
+
+        local_transfo_extra = GaussianBlur(p=0.5)
+
+        # normalization
+        self.normalize = v2.Compose(
+            [
+                v2.ToImage(),
+                v2.ToDtype(torch.float32, scale=True),
+                make_normalize_transform(mean=mean, std=std),
+            ]
+        )
+
+        if self.share_color_jitter:
+            self.color_jittering = color_jittering
+            self.global_transfo1 = v2.Compose([resize_global, global_transfo1_extra, self.normalize])
+            self.global_transfo2 = v2.Compose([resize_global, global_transfo2_extra, self.normalize])
+            self.local_transfo = v2.Compose([local_transfo_extra, self.normalize])
+        else:
+            self.global_transfo1 = v2.Compose(
+                [resize_global, color_jittering, global_transfo1_extra, self.normalize]
+            )
+            self.global_transfo2 = v2.Compose(
+                [resize_global, color_jittering, global_transfo2_extra, self.normalize]
+            )
+            self.local_transfo = v2.Compose([color_jittering, local_transfo_extra, self.normalize])
+
+    def __call__(self, image):
+        output = {}
+        output["weak_flag"] = True  # some residual from mugs
+
+        if self.share_color_jitter:
+            image = self.color_jittering(image)
+
+        # global crops:
+        im1_base = self.geometric_augmentation_global(image)
+        global_crop_1_transf = self.global_transfo1(im1_base)
+        global_crop_1 = self.resize_global_post_transf(global_crop_1_transf)
+
+        im2_base = self.geometric_augmentation_global(image)
+        global_crop_2_transf = self.global_transfo2(im2_base)
+        global_crop_2 = self.resize_global_post_transf(global_crop_2_transf)
+
+        output["global_crops"] = [global_crop_1, global_crop_2]
+
+        # global crops for teacher:
+        if self.teacher_no_color_jitter:
+            output["global_crops_teacher"] = [
+                self.normalize(im1_base),
+                self.normalize(im2_base),
+            ]
+        else:
+            output["global_crops_teacher"] = [global_crop_1, global_crop_2]
+
+        if self.gram_teacher_crops_size is not None:
+            # crops for gram teacher:
+            if self.gram_teacher_no_distortions:
+                gram_crop_1 = self.normalize(self.resize_gram_teacher(im1_base))
+                gram_crop_2 = self.normalize(self.resize_gram_teacher(im2_base))
+            else:
+                gram_crop_1 = self.resize_gram_teacher(global_crop_1_transf)
+                gram_crop_2 = self.resize_gram_teacher(global_crop_2_transf)
+            output["gram_teacher_crops"] = [gram_crop_1, gram_crop_2]
+
+        # local crops:
+        if self.local_crops_subset_of_global_crops:
+            _local_crops = [self.local_transfo(im1_base) for _ in range(self.local_crops_number // 2)] + [
+                self.local_transfo(im2_base) for _ in range(self.local_crops_number // 2)
+            ]
+
+            local_crops = []
+            offsets = []
+            gs = self.global_crops_size
+            ls = self.local_crops_size
+            for img in _local_crops:
+                rx, ry = np.random.randint(0, (gs - ls) // self.patch_size, 2) * self.patch_size
+                local_crops.append(img[:, rx : rx + ls, ry : ry + ls])
+                offsets.append((rx, ry))
+
+            output["local_crops"] = local_crops
+            output["offsets"] = offsets
+        else:
+            local_crops = [
+                self.local_transfo(self.geometric_augmentation_local(image)) for _ in range(self.local_crops_number)
+            ]
+            output["local_crops"] = local_crops
+            output["offsets"] = ()
+
+        return output

dinov3/data/collate.py

@@ -0,0 +1,125 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import random
+
+import torch
+
+
+def collate_data_and_cast(
+    samples_list,
+    mask_ratio_tuple,
+    mask_probability,
+    dtype,
+    n_tokens=None,
+    mask_generator=None,
+    random_circular_shift=False,
+    local_batch_size=None,
+):
+    n_global_crops = len(samples_list[0][0]["global_crops"])
+    n_local_crops = len(samples_list[0][0]["local_crops"])
+
+    collated_global_crops = torch.stack(
+        [s[0]["global_crops"][i] for i in range(n_global_crops) for s in samples_list]
+    )  # [n_global_crops, B, ...]
+    collated_local_crops = torch.stack([s[0]["local_crops"][i] for i in range(n_local_crops) for s in samples_list])
+    if "gram_teacher_crops" in samples_list[0][0]:
+        collated_gram_teacher_crops = torch.stack(
+            [s[0]["gram_teacher_crops"][i] for i in range(n_global_crops) for s in samples_list]
+        )  # [n_global_crops, B, ...]
+    else:
+        collated_gram_teacher_crops = None
+
+    if local_batch_size is not None:
+        # multi-distillation case, number of masks is different because the number of samples masked
+        # is different of the number of samples passed into the teacher initially
+        B = n_global_crops * local_batch_size
+    else:
+        B = len(collated_global_crops)
+    N = n_tokens
+    n_samples_masked = int(B * mask_probability)
+    probs = torch.linspace(*mask_ratio_tuple, n_samples_masked + 1)
+    upperbound = 0
+    masks_list = []
+    for i in range(0, n_samples_masked):
+        prob_max = probs[i + 1]
+        mask = torch.BoolTensor(mask_generator(int(N * prob_max)))
+        if random_circular_shift:  # apply le random circular shift to
+            shift_x, shift_y = (
+                random.randint(0, mask.shape[0] - 1),
+                random.randint(0, mask.shape[1] - 1),
+            )
+            mask = torch.roll(mask, (shift_x, shift_y), (0, 1))
+        masks_list.append(mask)
+        upperbound += int(N * prob_max)
+    for _ in range(n_samples_masked, B):
+        masks_list.append(torch.BoolTensor(mask_generator(0)))
+
+    random.shuffle(masks_list)
+
+    collated_masks = torch.stack(masks_list).flatten(1)
+    mask_indices_list = collated_masks.flatten().nonzero().flatten()
+
+    masks_weight = (1 / collated_masks.sum(-1).clamp(min=1.0)).unsqueeze(-1).expand_as(collated_masks)[collated_masks]
+
+    out = {
+        "collated_global_crops": collated_global_crops.to(dtype),
+        "collated_local_crops": collated_local_crops.to(dtype),
+        "collated_masks": collated_masks,
+        "mask_indices_list": mask_indices_list,
+        "masks_weight": masks_weight,
+        "upperbound": upperbound,
+        "n_masked_patches": torch.full((1,), fill_value=mask_indices_list.shape[0], dtype=torch.long),
+    }
+    if collated_gram_teacher_crops is not None:
+        out["collated_gram_teacher_crops"] = collated_gram_teacher_crops.to(dtype)
+    return out
+
+
+# def get_batch_subset(collated_data_batch, target_bs):
+def get_batch_subset(collated_data_batch, divide_by):
+    old_bs = collated_data_batch["collated_global_crops"].shape[0] // 2
+    target_bs = (old_bs + divide_by - 1) // divide_by
+    collated_global_crops = (
+        collated_data_batch["collated_global_crops"].unflatten(0, (2, old_bs)).narrow(1, 0, target_bs).flatten(0, 1)
+    )
+    collated_local_crops = (
+        collated_data_batch["collated_local_crops"].unflatten(0, (-1, old_bs)).narrow(1, 0, target_bs).flatten(0, 1)
+    )
+
+    masks_old_bs = collated_data_batch["collated_masks"].shape[0] // 2
+    masks_target_bs = masks_old_bs // divide_by
+    collated_masks = (
+        collated_data_batch["collated_masks"]
+        .unflatten(0, (2, masks_old_bs))
+        .narrow(1, 0, masks_target_bs)
+        .flatten(0, 1)
+    )
+    mask_indices_list = collated_masks.flatten().nonzero().flatten()
+
+    while mask_indices_list.shape[0] == 0:
+        _unbind = list(collated_data_batch["collated_masks"].unbind(0))
+        random.shuffle(_unbind)
+        _bind = torch.stack(_unbind, dim=0)
+        collated_masks = _bind.unflatten(0, (2, masks_old_bs)).narrow(1, 0, masks_target_bs).flatten(0, 1)
+        mask_indices_list = collated_masks.flatten().nonzero().flatten()
+
+    masks_weight = (1 / collated_masks.sum(-1).clamp(min=1.0)).unsqueeze(-1).expand_as(collated_masks)[collated_masks]
+    upperbound = collated_data_batch["upperbound"]
+
+    new_batch = {
+        "collated_global_crops": collated_global_crops,
+        "collated_local_crops": collated_local_crops,
+        "collated_masks": collated_masks,
+        "mask_indices_list": mask_indices_list,
+        "masks_weight": masks_weight,
+        "upperbound": upperbound,
+        "n_masked_patches": torch.full((1,), fill_value=mask_indices_list.shape[0], dtype=torch.long),
+    }
+
+    if "global_batch_size" in collated_data_batch.keys():
+        new_batch["global_batch_size"] = collated_data_batch["global_batch_size"] // divide_by
+
+    return new_batch

dinov3/data/datasets/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from .ade20k import ADE20K
+from .coco_captions import CocoCaptions
+from .image_net import ImageNet
+from .image_net_22k import ImageNet22k

dinov3/data/datasets/ade20k.py

@@ -0,0 +1,88 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import os
+from enum import Enum
+from typing import Any, Callable, List, Optional, Tuple, Union
+
+from PIL import Image
+
+from .decoders import Decoder, DenseTargetDecoder, ImageDataDecoder
+from .extended import ExtendedVisionDataset
+
+
+class _Split(Enum):
+    TRAIN = "train"
+    VAL = "val"
+
+    @property
+    def dirname(self) -> str:
+        return {
+            _Split.TRAIN: "training",
+            _Split.VAL: "validation",
+        }[self]
+
+
+def _file_to_segmentation_path(file_name: str, segm_base_path: str) -> str:
+    file_name_noext = os.path.splitext(file_name)[0]
+    return os.path.join(segm_base_path, file_name_noext + ".png")
+
+
+def _load_segmentation(root: str, split_file_names: List[str]):
+    segm_base_path = "annotations"
+    segmentation_paths = [_file_to_segmentation_path(file_name, segm_base_path) for file_name in split_file_names]
+    return segmentation_paths
+
+
+def _load_file_paths(root: str, split: _Split) -> Tuple[List[str], List[str]]:
+    with open(os.path.join(root, f"ADE20K_object150_{split.value}.txt")) as f:
+        split_file_names = sorted(f.read().strip().split("\n"))
+
+    all_segmentation_paths = _load_segmentation(root, split_file_names)
+    file_names = [os.path.join("images", el) for el in split_file_names]
+    return file_names, all_segmentation_paths
+
+
+class ADE20K(ExtendedVisionDataset):
+    Split = Union[_Split]
+    Labels = Union[Image.Image]
+
+    def __init__(
+        self,
+        split: "ADE20K.Split",
+        root: Optional[str] = None,
+        transforms: Optional[Callable] = None,
+        transform: Optional[Callable] = None,
+        target_transform: Optional[Callable] = None,
+        image_decoder: Decoder = ImageDataDecoder,
+        target_decoder: Decoder = DenseTargetDecoder,
+    ) -> None:
+        super().__init__(
+            root=root,
+            transforms=transforms,
+            transform=transform,
+            target_transform=target_transform,
+            image_decoder=image_decoder,
+            target_decoder=target_decoder,
+        )
+
+        self.image_paths, self.target_paths = _load_file_paths(root, split)
+
+    def get_image_data(self, index: int) -> bytes:
+        image_relpath = self.image_paths[index]
+        image_full_path = os.path.join(self.root, image_relpath)
+        with open(image_full_path, mode="rb") as f:
+            image_data = f.read()
+        return image_data
+
+    def get_target(self, index: int) -> Any:
+        target_relpath = self.target_paths[index]
+        target_full_path = os.path.join(self.root, target_relpath)
+        with open(target_full_path, mode="rb") as f:
+            target_data = f.read()
+        return target_data
+
+    def __len__(self) -> int:
+        return len(self.image_paths)

dinov3/data/datasets/coco_captions.py

@@ -0,0 +1,88 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import json
+import os
+import random
+from enum import Enum
+from typing import Callable, Dict, List, Optional, Union
+
+from .decoders import ImageDataDecoder, TargetDecoder
+from .extended import ExtendedVisionDataset
+
+# Dataset: https://www.kaggle.com/datasets/nikhil7280/coco-image-caption
+
+
+class _Split(Enum):
+    TRAIN = "train"
+    VAL = "val"
+
+
+def read_images_and_captions(root: str, split: _Split) -> List[Dict]:
+    image_dir = None
+    if _Split(split) == _Split.TRAIN:
+        annotations_full_path = os.path.join(
+            root, "annotations_trainval2014/annotations/captions_train2014.json"
+        )
+        image_dir = os.path.join(root, "train2014/train2014")
+    else:
+        annotations_full_path = os.path.join(
+            root, "annotations_trainval2017/annotations/captions_train2017.json"
+        )
+        image_dir = os.path.join(root, "val2017/val2017")
+    with open(annotations_full_path) as f:
+        all_annotations = json.load(f)
+    data = {}
+    for item in all_annotations["images"]:
+        id = item["id"]
+        data[id] = {
+            "id": None,
+            "image": os.path.join(image_dir, item["file_name"]),
+            "captions": [],
+        }
+    for item in all_annotations["annotations"]:
+        data[item["image_id"]]["id"] = item["image_id"]
+        data[item["image_id"]]["captions"].append(item["caption"])
+    return list(data.values())
+
+
+class CocoCaptions(ExtendedVisionDataset):
+    Split = Union[_Split]
+
+    def __init__(
+        self,
+        *,
+        split: "CocoCaptions.Split",
+        root: Optional[str] = None,
+        transforms: Optional[Callable] = None,
+        transform: Optional[Callable] = None,
+        target_transform: Optional[Callable] = None,
+    ) -> None:
+        super().__init__(
+            root=root,
+            transforms=transforms,
+            transform=transform,
+            target_transform=target_transform,
+            image_decoder=ImageDataDecoder,
+            target_decoder=TargetDecoder,
+        )
+
+        self.image_captions = read_images_and_captions(root, split)
+
+    def get_image_relpath(self, index: int) -> str:
+        image_path = self.image_captions[index]["image"]
+        return image_path
+
+    def get_image_data(self, index: int) -> bytes:
+        image_path = self.get_image_relpath(index)
+        with open(image_path, mode="rb") as f:
+            image_data = f.read()
+        return image_data
+
+    def get_target(self, index: int) -> str:
+        return random.choice(self.image_captions[index]["captions"])
+
+    def __len__(self) -> int:
+        return len(self.image_captions)

dinov3/data/datasets/decoders.py

@@ -0,0 +1,40 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from io import BytesIO
+from typing import Any
+
+from PIL import Image
+
+
+class Decoder:
+    def decode(self) -> Any:
+        raise NotImplementedError
+
+
+class ImageDataDecoder(Decoder):
+    def __init__(self, image_data: bytes) -> None:
+        self._image_data = image_data
+
+    def decode(self) -> Image:
+        f = BytesIO(self._image_data)
+        return Image.open(f).convert(mode="RGB")
+
+
+class TargetDecoder(Decoder):
+    def __init__(self, target: Any):
+        self._target = target
+
+    def decode(self) -> Any:
+        return self._target
+
+
+class DenseTargetDecoder(Decoder):
+    def __init__(self, image_data: bytes) -> None:
+        self._image_data = image_data
+
+    def decode(self) -> Image:
+        f = BytesIO(self._image_data)
+        return Image.open(f)

dinov3/data/datasets/extended.py

@@ -0,0 +1,46 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from typing import Any, Tuple
+
+from torchvision.datasets import VisionDataset
+
+from .decoders import Decoder, ImageDataDecoder, TargetDecoder
+
+
+class ExtendedVisionDataset(VisionDataset):
+    def __init__(
+        self,
+        image_decoder: Decoder = ImageDataDecoder,
+        target_decoder: Decoder = TargetDecoder,
+        *args,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)  # type: ignore
+        self.image_decoder = image_decoder
+        self.target_decoder = target_decoder
+
+    def get_image_data(self, index: int) -> bytes:
+        raise NotImplementedError
+
+    def get_target(self, index: int) -> Any:
+        raise NotImplementedError
+
+    def __getitem__(self, index: int) -> Tuple[Any, Any]:
+        try:
+            image_data = self.get_image_data(index)
+            image = self.image_decoder(image_data).decode()
+        except Exception as e:
+            raise RuntimeError(f"can not read image for sample {index}") from e
+        target = self.get_target(index)
+        target = self.target_decoder(target).decode()
+
+        if self.transforms is not None:
+            image, target = self.transforms(image, target)
+
+        return image, target
+
+    def __len__(self) -> int:
+        raise NotImplementedError

dinov3/data/datasets/image_net.py

@@ -0,0 +1,297 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import csv
+import logging
+import os
+from enum import Enum
+from typing import Callable, List, Optional, Tuple, Union
+
+import numpy as np
+
+from .decoders import ImageDataDecoder, TargetDecoder
+from .extended import ExtendedVisionDataset
+
+logger = logging.getLogger("dinov3")
+_Target = int
+
+
+class _Split(Enum):
+    TRAIN = "train"
+    VAL = "val"
+    TEST = "test"  # NOTE: torchvision does not support the test split
+
+    @property
+    def length(self) -> int:
+        split_lengths = {
+            _Split.TRAIN: 1_281_167,
+            _Split.VAL: 50_000,
+            _Split.TEST: 100_000,
+        }
+        return split_lengths[self]
+
+    def get_dirname(self, class_id: Optional[str] = None) -> str:
+        return self.value if class_id is None else os.path.join(self.value, class_id)
+
+    def get_image_relpath(self, actual_index: int, class_id: Optional[str] = None) -> str:
+        dirname = self.get_dirname(class_id)
+        if self == _Split.TRAIN:
+            basename = f"{class_id}_{actual_index}"
+        else:  # self in (_Split.VAL, _Split.TEST):
+            basename = f"ILSVRC2012_{self.value}_{actual_index:08d}"
+        return os.path.join(dirname, basename + ".JPEG")
+
+    def parse_image_relpath(self, image_relpath: str) -> Tuple[str, int]:
+        assert self != _Split.TEST
+        dirname, filename = os.path.split(image_relpath)
+        class_id = os.path.split(dirname)[-1]
+        basename, _ = os.path.splitext(filename)
+        actual_index = int(basename.split("_")[-1])
+        return class_id, actual_index
+
+
+class ImageNet(ExtendedVisionDataset):
+    Target = Union[_Target]
+    Split = Union[_Split]
+
+    def __init__(
+        self,
+        *,
+        split: "ImageNet.Split",
+        root: str,
+        extra: str,
+        transforms: Optional[Callable] = None,
+        transform: Optional[Callable] = None,
+        target_transform: Optional[Callable] = None,
+    ) -> None:
+        super().__init__(
+            root=root,
+            transforms=transforms,
+            transform=transform,
+            target_transform=target_transform,
+            image_decoder=ImageDataDecoder,
+            target_decoder=TargetDecoder,
+        )
+        self._extra_root = extra
+        self._split = split
+
+        self._entries = None
+        self._class_ids = None
+        self._class_names = None
+
+    @property
+    def split(self) -> "ImageNet.Split":
+        return self._split
+
+    def _get_extra_full_path(self, extra_path: str) -> str:
+        return os.path.join(self._extra_root, extra_path)
+
+    def _load_extra(self, extra_path: str) -> np.ndarray:
+        extra_full_path = self._get_extra_full_path(extra_path)
+        return np.load(extra_full_path, mmap_mode="r")
+
+    def _save_extra(self, extra_array: np.ndarray, extra_path: str) -> None:
+        extra_full_path = self._get_extra_full_path(extra_path)
+        os.makedirs(self._extra_root, exist_ok=True)
+        np.save(extra_full_path, extra_array)
+
+    @property
+    def _entries_path(self) -> str:
+        return f"entries-{self._split.value.upper()}.npy"
+
+    @property
+    def _class_ids_path(self) -> str:
+        return f"class-ids-{self._split.value.upper()}.npy"
+
+    @property
+    def _class_names_path(self) -> str:
+        return f"class-names-{self._split.value.upper()}.npy"
+
+    def _get_entries(self) -> np.ndarray:
+        if self._entries is None:
+            self._entries = self._load_extra(self._entries_path)
+        assert self._entries is not None
+        return self._entries
+
+    def _get_class_ids(self) -> np.ndarray:
+        if self._split == _Split.TEST:
+            raise AssertionError("Class IDs are not available in TEST split")
+        if self._class_ids is None:
+            self._class_ids = self._load_extra(self._class_ids_path)
+        assert self._class_ids is not None
+        return self._class_ids
+
+    def _get_class_names(self) -> np.ndarray:
+        if self._split == _Split.TEST:
+            raise AssertionError("Class names are not available in TEST split")
+        if self._class_names is None:
+            self._class_names = self._load_extra(self._class_names_path)
+        assert self._class_names is not None
+        return self._class_names
+
+    def find_class_id(self, class_index: int) -> str:
+        class_ids = self._get_class_ids()
+        return str(class_ids[class_index])
+
+    def find_class_name(self, class_index: int) -> str:
+        class_names = self._get_class_names()
+        return str(class_names[class_index])
+
+    def get_image_data(self, index: int) -> bytes:
+        entries = self._get_entries()
+        actual_index = entries[index]["actual_index"]
+
+        class_id = self.get_class_id(index)
+
+        image_relpath = self.split.get_image_relpath(actual_index, class_id)
+        image_full_path = os.path.join(self.root, image_relpath)
+        with open(image_full_path, mode="rb") as f:
+            image_data = f.read()
+        return image_data
+
+    def get_target(self, index: int) -> Optional[Target]:
+        entries = self._get_entries()
+        class_index = entries[index]["class_index"]
+        return None if self.split == _Split.TEST else int(class_index)
+
+    def get_targets(self) -> Optional[np.ndarray]:
+        entries = self._get_entries()
+        return None if self.split == _Split.TEST else entries["class_index"]
+
+    def get_class_id(self, index: int) -> Optional[str]:
+        entries = self._get_entries()
+        class_id = entries[index]["class_id"]
+        return None if self.split == _Split.TEST else str(class_id)
+
+    def get_class_name(self, index: int) -> Optional[str]:
+        entries = self._get_entries()
+        class_name = entries[index]["class_name"]
+        return None if self.split == _Split.TEST else str(class_name)
+
+    def __len__(self) -> int:
+        entries = self._get_entries()
+        assert len(entries) == self.split.length
+        return len(entries)
+
+    def _load_labels(self, labels_path: str) -> List[Tuple[str, str]]:
+        labels_full_path = os.path.join(self.root, labels_path)
+        labels = []
+
+        try:
+            with open(labels_full_path, "r") as f:
+                reader = csv.reader(f)
+                for row in reader:
+                    class_id, class_name = row
+                    labels.append((class_id, class_name))
+        except OSError as e:
+            raise RuntimeError(f'can not read labels file "{labels_full_path}"') from e
+
+        return labels
+
+    def _dump_entries(self) -> None:
+        split = self.split
+        if split == ImageNet.Split.TEST:
+            dataset = None
+            sample_count = split.length
+            max_class_id_length, max_class_name_length = 0, 0
+        else:
+            labels_path = "labels.txt"
+            logger.info(f'loading labels from "{labels_path}"')
+            labels = self._load_labels(labels_path)
+
+            # NOTE: Using torchvision ImageFolder for consistency
+            from torchvision.datasets import ImageFolder
+
+            dataset_root = os.path.join(self.root, split.get_dirname())
+            dataset = ImageFolder(dataset_root)
+            sample_count = len(dataset)
+            max_class_id_length, max_class_name_length = -1, -1
+            for sample in dataset.samples:
+                _, class_index = sample
+                class_id, class_name = labels[class_index]
+                max_class_id_length = max(len(class_id), max_class_id_length)
+                max_class_name_length = max(len(class_name), max_class_name_length)
+
+        dtype = np.dtype(
+            [
+                ("actual_index", "<u4"),
+                ("class_index", "<u4"),
+                ("class_id", f"U{max_class_id_length}"),
+                ("class_name", f"U{max_class_name_length}"),
+            ]
+        )
+        entries_array = np.empty(sample_count, dtype=dtype)
+
+        if split == ImageNet.Split.TEST:
+            old_percent = -1
+            for index in range(sample_count):
+                percent = 100 * (index + 1) // sample_count
+                if percent > old_percent:
+                    logger.info(f"creating entries: {percent}%")
+                    old_percent = percent
+
+                actual_index = index + 1
+                class_index = np.uint32(-1)
+                class_id, class_name = "", ""
+                entries_array[index] = (actual_index, class_index, class_id, class_name)
+        else:
+            class_names = {class_id: class_name for class_id, class_name in labels}
+
+            assert dataset
+            old_percent = -1
+            for index in range(sample_count):
+                percent = 100 * (index + 1) // sample_count
+                if percent > old_percent:
+                    logger.info(f"creating entries: {percent}%")
+                    old_percent = percent
+
+                image_full_path, class_index = dataset.samples[index]
+                image_relpath = os.path.relpath(image_full_path, self.root)
+                class_id, actual_index = split.parse_image_relpath(image_relpath)
+                class_name = class_names[class_id]
+                entries_array[index] = (actual_index, class_index, class_id, class_name)
+
+        logger.info(f'saving entries to "{self._entries_path}"')
+        self._save_extra(entries_array, self._entries_path)
+
+    def _dump_class_ids_and_names(self) -> None:
+        split = self.split
+        if split == ImageNet.Split.TEST:
+            return
+
+        entries_array = self._load_extra(self._entries_path)
+
+        max_class_id_length, max_class_name_length, max_class_index = -1, -1, -1
+        for entry in entries_array:
+            class_index, class_id, class_name = (
+                entry["class_index"],
+                entry["class_id"],
+                entry["class_name"],
+            )
+            max_class_index = max(int(class_index), max_class_index)
+            max_class_id_length = max(len(str(class_id)), max_class_id_length)
+            max_class_name_length = max(len(str(class_name)), max_class_name_length)
+
+        class_count = max_class_index + 1
+        class_ids_array = np.empty(class_count, dtype=f"U{max_class_id_length}")
+        class_names_array = np.empty(class_count, dtype=f"U{max_class_name_length}")
+        for entry in entries_array:
+            class_index, class_id, class_name = (
+                entry["class_index"],
+                entry["class_id"],
+                entry["class_name"],
+            )
+            class_ids_array[class_index] = class_id
+            class_names_array[class_index] = class_name
+
+        logger.info(f'saving class IDs to "{self._class_ids_path}"')
+        self._save_extra(class_ids_array, self._class_ids_path)
+
+        logger.info(f'saving class names to "{self._class_names_path}"')
+        self._save_extra(class_names_array, self._class_names_path)
+
+    def dump_extra(self) -> None:
+        self._dump_entries()
+        self._dump_class_ids_and_names()

dinov3/data/datasets/image_net_22k.py

@@ -0,0 +1,301 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import os
+import warnings
+from dataclasses import dataclass
+from enum import Enum
+from functools import lru_cache
+from gzip import GzipFile
+from io import BytesIO
+from mmap import ACCESS_READ, mmap
+from typing import Any, Callable, List, Optional, Set, Tuple
+
+import numpy as np
+
+from .extended import ExtendedVisionDataset
+
+_Labels = int
+
+_DEFAULT_MMAP_CACHE_SIZE = 16  # Warning: This can exhaust file descriptors
+
+
+@dataclass
+class _ClassEntry:
+    block_offset: int
+    maybe_filename: Optional[str] = None
+
+
+@dataclass
+class _Entry:
+    class_index: int  # noqa: E701
+    start_offset: int
+    end_offset: int
+    filename: str
+
+
+class _Split(Enum):
+    TRAIN = "train"
+    VAL = "val"
+
+    @property
+    def length(self) -> int:
+        return {
+            _Split.TRAIN: 11_797_647,
+            _Split.VAL: 561_050,
+        }[self]
+
+    def entries_path(self):
+        return f"imagenet21kp_{self.value}.txt"
+
+
+def _get_tarball_path(class_id: str) -> str:
+    return f"{class_id}.tar"
+
+
+def _make_mmap_tarball(tarballs_root: str, mmap_cache_size: int):
+    @lru_cache(maxsize=mmap_cache_size)
+    def _mmap_tarball(class_id: str) -> mmap:
+        tarball_path = _get_tarball_path(class_id)
+        tarball_full_path = os.path.join(tarballs_root, tarball_path)
+        with open(tarball_full_path) as f:
+            return mmap(fileno=f.fileno(), length=0, access=ACCESS_READ)
+
+    return _mmap_tarball
+
+
+class ImageNet22k(ExtendedVisionDataset):
+    _GZIPPED_INDICES: Set[int] = {
+        841_545,
+        1_304_131,
+        2_437_921,
+        2_672_079,
+        2_795_676,
+        2_969_786,
+        6_902_965,
+        6_903_550,
+        6_903_628,
+        7_432_557,
+        7_432_589,
+        7_813_809,
+        8_329_633,
+        10_296_990,
+        10_417_652,
+        10_492_265,
+        10_598_078,
+        10_782_398,
+        10_902_612,
+        11_203_736,
+        11_342_890,
+        11_397_596,
+        11_589_762,
+        11_705_103,
+        12_936_875,
+        13_289_782,
+    }
+    Labels = _Labels
+
+    def __init__(
+        self,
+        *,
+        root: str,
+        extra: str,
+        transforms: Optional[Callable] = None,
+        transform: Optional[Callable] = None,
+        target_transform: Optional[Callable] = None,
+        mmap_cache_size: int = _DEFAULT_MMAP_CACHE_SIZE,
+    ) -> None:
+        super().__init__(root, transforms, transform, target_transform)
+        self._extra_root = extra
+
+        entries_path = self._get_entries_path(root)
+        self._entries = self._load_extra(entries_path)
+
+        class_ids_path = self._get_class_ids_path(root)
+        self._class_ids = self._load_extra(class_ids_path)
+
+        self._gzipped_indices = ImageNet22k._GZIPPED_INDICES
+        self._mmap_tarball = _make_mmap_tarball(self._tarballs_root, mmap_cache_size)
+
+    def _get_entries_path(self, root: Optional[str] = None) -> str:
+        return "entries.npy"
+
+    def _get_class_ids_path(self, root: Optional[str] = None) -> str:
+        return "class-ids.npy"
+
+    def _find_class_ids(self, path: str) -> List[str]:
+        class_ids = []
+
+        with os.scandir(path) as entries:
+            for entry in entries:
+                root, ext = os.path.splitext(entry.name)
+                if ext != ".tar":
+                    continue
+                class_ids.append(root)
+
+        return sorted(class_ids)
+
+    def _load_entries_class_ids(self, root: Optional[str] = None) -> Tuple[List[_Entry], List[str]]:
+        root = self.get_root(root)
+        entries: List[_Entry] = []
+        class_ids = self._find_class_ids(root)
+
+        for class_index, class_id in enumerate(class_ids):
+            path = os.path.join(root, "blocks", f"{class_id}.log")
+            class_entries = []
+
+            try:
+                with open(path) as f:
+                    for line in f:
+                        line = line.rstrip()
+                        block, filename = line.split(":")
+                        block_offset = int(block[6:])
+                        filename = filename[1:]
+
+                        maybe_filename = None
+                        if filename != "** Block of NULs **":
+                            maybe_filename = filename
+                            _, ext = os.path.splitext(filename)
+                            # assert ext == ".JPEG"
+
+                        class_entry = _ClassEntry(block_offset, maybe_filename)
+                        class_entries.append(class_entry)
+            except OSError as e:
+                raise RuntimeError(f'can not read blocks file "{path}"') from e
+
+            assert class_entries[-1].maybe_filename is None
+
+            for class_entry1, class_entry2 in zip(class_entries, class_entries[1:]):
+                assert class_entry1.block_offset <= class_entry2.block_offset
+                start_offset = 512 * class_entry1.block_offset
+                end_offset = 512 * class_entry2.block_offset
+                assert class_entry1.maybe_filename is not None
+                filename = class_entry1.maybe_filename
+                entry = _Entry(class_index, start_offset, end_offset, filename)
+                # Skip invalid image files (PIL throws UnidentifiedImageError)
+                if filename == "n06470073_47249.JPEG":
+                    continue
+                entries.append(entry)
+
+        return entries, class_ids
+
+    def _load_extra(self, extra_path: str) -> np.ndarray:
+        extra_root = self._extra_root
+        extra_full_path = os.path.join(extra_root, extra_path)
+        return np.load(extra_full_path, mmap_mode="r")
+
+    def _save_extra(self, extra_array: np.ndarray, extra_path: str) -> None:
+        extra_root = self._extra_root
+        extra_full_path = os.path.join(extra_root, extra_path)
+        os.makedirs(extra_root, exist_ok=True)
+        np.save(extra_full_path, extra_array)
+
+    @property
+    def _tarballs_root(self) -> str:
+        return self.root
+
+    def find_class_id(self, class_index: int) -> str:
+        return str(self._class_ids[class_index])
+
+    def get_image_data(self, index: int) -> bytes:
+        entry = self._entries[index]
+        class_id = entry["class_id"]
+        class_mmap = self._mmap_tarball(class_id)
+
+        start_offset, end_offset = entry["start_offset"], entry["end_offset"]
+        try:
+            mapped_data = class_mmap[start_offset:end_offset]
+            data = mapped_data[512:]  # Skip entry header block
+
+            if len(data) >= 2 and tuple(data[:2]) == (0x1F, 0x8B):
+                assert index in self._gzipped_indices, f"unexpected gzip header for sample {index}"
+                with GzipFile(fileobj=BytesIO(data)) as g:
+                    data = g.read()
+        except Exception as e:
+            raise RuntimeError(f"can not retrieve image data for sample {index} " f'from "{class_id}" tarball') from e
+
+        return data
+
+    def get_target(self, index: int) -> Any:
+        return int(self._entries[index]["class_index"])
+
+    def get_targets(self) -> np.ndarray:
+        return self._entries["class_index"]
+
+    def get_class_id(self, index: int) -> str:
+        return str(self._entries[index]["class_id"])
+
+    def get_class_ids(self) -> np.ndarray:
+        return self._entries["class_id"]
+
+    def __getitem__(self, index: int) -> Tuple[Any, Any]:
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            return super().__getitem__(index)
+
+    def __len__(self) -> int:
+        return len(self._entries)
+
+    def _dump_entries(self, *args, **kwargs) -> None:
+        entries, class_ids = self._load_entries_class_ids(*args, **kwargs)
+
+        max_class_id_length, max_filename_length, max_class_index = -1, -1, -1
+        for entry in entries:
+            class_id = class_ids[entry.class_index]
+            max_class_index = max(entry.class_index, max_class_index)
+            max_class_id_length = max(len(class_id), max_class_id_length)
+            max_filename_length = max(len(entry.filename), max_filename_length)
+
+        dtype = np.dtype(
+            [
+                ("class_index", "<u4"),
+                ("class_id", f"U{max_class_id_length}"),
+                ("start_offset", "<u4"),
+                ("end_offset", "<u4"),
+                ("filename", f"U{max_filename_length}"),
+            ]
+        )
+        sample_count = len(entries)
+        entries_array = np.empty(sample_count, dtype=dtype)
+        for i, entry in enumerate(entries):
+            class_index = entry.class_index
+            class_id = class_ids[class_index]
+            start_offset = entry.start_offset
+            end_offset = entry.end_offset
+            filename = entry.filename
+            entries_array[i] = (
+                class_index,
+                class_id,
+                start_offset,
+                end_offset,
+                filename,
+            )
+
+        entries_path = self._get_entries_path(*args, **kwargs)
+        self._save_extra(entries_array, entries_path)
+
+    def _dump_class_ids(self, *args, **kwargs) -> None:
+        entries_path = self._get_entries_path(*args, **kwargs)
+        entries_array = self._load_extra(entries_path)
+
+        max_class_id_length, max_class_index = -1, -1
+        for entry in entries_array:
+            class_index, class_id = entry["class_index"], entry["class_id"]
+            max_class_index = max(int(class_index), max_class_index)
+            max_class_id_length = max(len(str(class_id)), max_class_id_length)
+
+        class_ids_array = np.empty(max_class_index + 1, dtype=f"U{max_class_id_length}")
+        for entry in entries_array:
+            class_index, class_id = entry["class_index"], entry["class_id"]
+            class_ids_array[class_index] = class_id
+        class_ids_path = self._get_class_ids_path(*args, **kwargs)
+        self._save_extra(class_ids_array, class_ids_path)
+
+    def _dump_extra(self, *args, **kwargs) -> None:
+        self._dump_entries(*args, *kwargs)
+        self._dump_class_ids(*args, *kwargs)
+
+    def dump_extra(self, root: Optional[str] = None) -> None:
+        return self._dump_extra(root)

dinov3/data/loaders.py

@@ -0,0 +1,237 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import logging
+from enum import Enum
+from typing import Any, Callable, List, Optional, TypeVar
+
+import torch
+from torch.utils.data import Sampler
+
+from .datasets import ADE20K, CocoCaptions, ImageNet, ImageNet22k
+from .samplers import EpochSampler, InfiniteSampler, ShardedInfiniteSampler
+
+logger = logging.getLogger("dinov3")
+
+
+class SamplerType(Enum):
+    DISTRIBUTED = 0
+    EPOCH = 1
+    INFINITE = 2
+    SHARDED_INFINITE = 3
+    SHARDED_INFINITE_NEW = 4
+
+
+def _make_bool_str(b: bool) -> str:
+    return "yes" if b else "no"
+
+
+def _make_sample_transform(
+    image_transform: Optional[Callable] = None,
+    target_transform: Optional[Callable] = None,
+):
+    def transform(sample):
+        image, target = sample
+        if image_transform is not None:
+            image = image_transform(image)
+        if target_transform is not None:
+            target = target_transform(target)
+        return image, target
+
+    return transform
+
+
+def _parse_dataset_str(dataset_str: str):
+    tokens = dataset_str.split(":")
+
+    name = tokens[0]
+    kwargs = {}
+
+    for token in tokens[1:]:
+        key, value = token.split("=")
+        assert key in ("root", "extra", "split")
+        kwargs[key] = value
+
+    if name == "ImageNet":
+        class_ = ImageNet
+        if "split" in kwargs:
+            kwargs["split"] = ImageNet.Split[kwargs["split"]]
+    elif name == "ImageNet22k":
+        class_ = ImageNet22k
+    elif name == "ADE20K":
+        class_ = ADE20K
+        if "split" in kwargs:
+            kwargs["split"] = ADE20K.Split[kwargs["split"]]
+    elif name == "CocoCaptions":
+        class_ = CocoCaptions
+        if "split" in kwargs:
+            kwargs["split"] = CocoCaptions.Split[kwargs["split"]]
+    else:
+        raise ValueError(f'Unsupported dataset "{name}"')
+
+    return class_, kwargs
+
+
+def make_dataset(
+    *,
+    dataset_str: str,
+    transform: Optional[Callable] = None,
+    target_transform: Optional[Callable] = None,
+    transforms: Optional[Callable] = None,
+):
+    """
+    Creates a dataset with the specified parameters.
+
+    Args:
+        dataset_str: A dataset string description (e.g. ImageNet:split=TRAIN).
+        transform: A transform to apply to images.
+        target_transform: A transform to apply to targets.
+        transforms: A transform to apply to both images and targets.
+
+    Returns:
+        The created dataset.
+    """
+    logger.info(f'using dataset: "{dataset_str}"')
+
+    class_, kwargs = _parse_dataset_str(dataset_str)
+    dataset = class_(transform=transform, target_transform=target_transform, transforms=transforms, **kwargs)
+
+    logger.info(f"# of dataset samples: {len(dataset):,d}")
+
+    # Aggregated datasets do not expose (yet) these attributes, so add them.
+    if not hasattr(dataset, "transform"):
+        dataset.transform = transform
+    if not hasattr(dataset, "target_transform"):
+        dataset.target_transform = target_transform
+    if not hasattr(dataset, "transforms"):
+        dataset.transforms = transforms
+
+    return dataset
+
+
+def _make_sampler(
+    *,
+    dataset,
+    type: Optional[SamplerType] = None,
+    shuffle: bool = False,
+    seed: int = 0,
+    size: int = -1,
+    advance: int = 0,
+) -> Optional[Sampler]:
+    sample_count = len(dataset)
+
+    if type == SamplerType.INFINITE:
+        logger.info("sampler: infinite")
+        if size > 0:
+            raise ValueError("sampler size > 0 is invalid")
+        return InfiniteSampler(
+            sample_count=sample_count,
+            shuffle=shuffle,
+            seed=seed,
+            advance=advance,
+        )
+    elif type in (SamplerType.SHARDED_INFINITE, SamplerType.SHARDED_INFINITE_NEW):
+        logger.info("sampler: sharded infinite")
+        if size > 0:
+            raise ValueError("sampler size > 0 is invalid")
+        use_new_shuffle_tensor_slice = type == SamplerType.SHARDED_INFINITE_NEW
+        return ShardedInfiniteSampler(
+            sample_count=sample_count,
+            shuffle=shuffle,
+            seed=seed,
+            advance=advance,
+            use_new_shuffle_tensor_slice=use_new_shuffle_tensor_slice,
+        )
+    elif type == SamplerType.EPOCH:
+        logger.info("sampler: epoch")
+        if advance > 0:
+            raise NotImplementedError("sampler advance > 0 is not supported")
+        size = size if size > 0 else sample_count
+        logger.info(f"# of samples / epoch: {size:,d}")
+        return EpochSampler(
+            size=size,
+            sample_count=sample_count,
+            shuffle=shuffle,
+            seed=seed,
+        )
+    elif type == SamplerType.DISTRIBUTED:
+        logger.info("sampler: distributed")
+        if size > 0:
+            raise ValueError("sampler size > 0 is invalid")
+        if advance > 0:
+            raise ValueError("sampler advance > 0 is invalid")
+        return torch.utils.data.DistributedSampler(
+            dataset=dataset,
+            shuffle=shuffle,
+            seed=seed,
+            drop_last=False,
+        )
+
+    logger.info("sampler: none")
+    return None
+
+
+T = TypeVar("T")
+
+
+def make_data_loader(
+    *,
+    dataset,
+    batch_size: int,
+    num_workers: int,
+    shuffle: bool = True,
+    seed: int = 0,
+    sampler_type: Optional[SamplerType] = SamplerType.INFINITE,
+    sampler_size: int = -1,
+    sampler_advance: int = 0,
+    drop_last: bool = True,
+    persistent_workers: bool = False,
+    collate_fn: Optional[Callable[[List[T]], Any]] = None,
+    worker_init_fn: Optional[Callable[[List[T]], Any]] = None,
+):
+    """
+    Creates a data loader with the specified parameters.
+
+    Args:
+        dataset: A dataset (third party, LaViDa or WebDataset).
+        batch_size: The size of batches to generate.
+        num_workers: The number of workers to use.
+        shuffle: Whether to shuffle samples.
+        seed: The random seed to use.
+        sampler_type: Which sampler to use: EPOCH, INFINITE, SHARDED_INFINITE, SHARDED_INFINITE_NEW, DISTRIBUTED or None.
+        sampler_size: The number of images per epoch (when applicable) or -1 for the entire dataset.
+        sampler_advance: How many samples to skip (when applicable).
+        drop_last: Whether the last non-full batch of data should be dropped.
+        persistent_workers: maintain the workers Dataset instances alive after a dataset has been consumed once.
+        collate_fn: Function that performs batch collation
+    """
+
+    sampler = _make_sampler(
+        dataset=dataset,
+        type=sampler_type,
+        shuffle=shuffle,
+        seed=seed,
+        size=sampler_size,
+        advance=sampler_advance,
+    )
+
+    logger.info("using PyTorch data loader")
+    data_loader = torch.utils.data.DataLoader(
+        dataset,
+        sampler=sampler,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        pin_memory=True,
+        drop_last=drop_last,
+        persistent_workers=persistent_workers,
+        collate_fn=collate_fn,
+        worker_init_fn=worker_init_fn,
+    )
+
+    try:
+        logger.info(f"# of batches: {len(data_loader):,d}")
+    except TypeError:  # data loader has no length
+        logger.info("infinite data loader")
+    return data_loader

dinov3/data/masking.py

@@ -0,0 +1,94 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import math
+import random
+
+import numpy as np
+
+
+class MaskingGenerator:
+    def __init__(
+        self,
+        input_size,
+        num_masking_patches=None,
+        min_num_patches=4,
+        max_num_patches=None,
+        min_aspect=0.3,
+        max_aspect=None,
+    ):
+        if not isinstance(input_size, tuple):
+            input_size = (input_size,) * 2
+        self.height, self.width = input_size
+
+        self.num_patches = self.height * self.width
+        self.num_masking_patches = num_masking_patches
+
+        self.min_num_patches = min_num_patches
+        self.max_num_patches = num_masking_patches if max_num_patches is None else max_num_patches
+
+        max_aspect = max_aspect or 1 / min_aspect
+        self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect))
+
+    def __repr__(self):
+        repr_str = "Generator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" % (
+            self.height,
+            self.width,
+            self.min_num_patches,
+            self.max_num_patches,
+            self.num_masking_patches,
+            self.log_aspect_ratio[0],
+            self.log_aspect_ratio[1],
+        )
+        return repr_str
+
+    def get_shape(self):
+        return self.height, self.width
+
+    def _mask(self, mask, max_mask_patches):
+        delta = 0
+        for _ in range(10):
+            target_area = random.uniform(self.min_num_patches, max_mask_patches)
+            aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
+            h = int(round(math.sqrt(target_area * aspect_ratio)))
+            w = int(round(math.sqrt(target_area / aspect_ratio)))
+            if w < self.width and h < self.height:
+                top = random.randint(0, self.height - h)
+                left = random.randint(0, self.width - w)
+
+                num_masked = mask[top : top + h, left : left + w].sum()
+                # Overlap
+                if 0 < h * w - num_masked <= max_mask_patches:
+                    for i in range(top, top + h):
+                        for j in range(left, left + w):
+                            if mask[i, j] == 0:
+                                mask[i, j] = 1
+                                delta += 1
+
+                if delta > 0:
+                    break
+        return delta
+
+    def __call__(self, num_masking_patches=0):
+        mask = np.zeros(shape=self.get_shape(), dtype=bool)
+        mask_count = 0
+        while mask_count < num_masking_patches:
+            max_mask_patches = num_masking_patches - mask_count
+            max_mask_patches = min(max_mask_patches, self.max_num_patches)
+
+            delta = self._mask(mask, max_mask_patches)
+            if delta == 0:
+                break
+            else:
+                mask_count += delta
+
+        return self.complete_mask_randomly(mask, num_masking_patches)
+
+    def complete_mask_randomly(self, mask, num_masking_patches):
+        shape = mask.shape
+        m2 = mask.flatten()
+        to_add = np.random.choice(np.where(~m2)[0], size=num_masking_patches - m2.sum(), replace=False)
+        m2[to_add] = True
+        return m2.reshape(shape)

dinov3/data/meta_loaders.py

@@ -0,0 +1,78 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import logging
+from typing import Any, Iterable, Iterator, List, Tuple, TypeVar
+
+import numpy as np
+
+logger = logging.getLogger("dinov3")
+Loader = Iterable[List[Any]]
+T = TypeVar("T")
+
+
+class CombinedDataLoader:
+    """
+    Combines data loaders using the provided sampling ratios
+    """
+
+    GLOBAL_HOMOGENEOUS = 0
+    LOCAL_HOMOGENEOUS = 1
+
+    def __init__(
+        self,
+        loaders_with_ratios: Iterable[Tuple[Loader, float]],
+        batch_size: int,
+        combining_mode: int = 1,
+        seed: int = 65537,
+        name: str = None,
+        logging_period: int = 100,
+    ):
+        if combining_mode not in [self.GLOBAL_HOMOGENEOUS, self.LOCAL_HOMOGENEOUS]:
+            raise ValueError(f"Unsupported value of combining_mode ({combining_mode})")
+        loaders, ratios = zip(*loaders_with_ratios)
+        assert np.all([loader.batch_size == batch_size for loader in loaders]), (
+            f"All individual loaders must have the same batch size to the combined data loader for combining_mode={combining_mode}"
+        )
+        self.loaders = loaders
+        self.ratios = ratios
+        self.batch_size = batch_size
+        self.combining_mode = combining_mode
+        self.initial_seed = seed
+        self.name = name if name is not None else ""
+        self.logging_period = logging_period
+        if combining_mode == self.GLOBAL_HOMOGENEOUS:
+            logger.info(f"Initialize CDL {self.name} with seed={seed}")
+            self.seed = seed
+            self.rng = np.random.default_rng(seed=seed)
+        else:
+            logger.info(f"Initialize CDL {self.name} with random seed")
+            self.seed = 0
+            self.rng = np.random.default_rng()
+        self.loader_count = np.zeros(len(self.loaders))
+
+    def homogeneous_iterator(self) -> Iterator[List[Any]]:
+        iteration = 0
+        iters = [iter(loader) for loader in self.loaders]
+        while True:
+            iteration += 1
+            try:
+                idx = self.rng.choice(len(self.loaders), p=self.ratios)
+                self.loader_count[idx] += 1
+                if iteration % self.logging_period == 0:
+                    logger.info(f"Empirical ratios: CDL {self.name} {self.loader_count / self.loader_count.sum()}")
+                yield next(iters[idx])
+            except StopIteration:
+                break
+
+    def heterogeneous_iterator(self) -> Iterator[List[Any]]:
+        pass
+
+    def __iter__(self) -> Iterator[List[Any]]:
+        if self.combining_mode in [self.GLOBAL_HOMOGENEOUS, self.LOCAL_HOMOGENEOUS]:
+            logger.info(f"Using homogeneous iterator for CDL {self.name}")
+            return self.homogeneous_iterator()
+        else:
+            raise ValueError(f"Unsupported value of combining_mode ({self.combining_mode})")

dinov3/data/samplers.py

@@ -0,0 +1,229 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import itertools
+import warnings
+from typing import Any, Optional
+
+import numpy as np
+import torch
+from torch.utils.data.sampler import Sampler
+
+from dinov3.distributed import get_rank, get_world_size
+
+
+class EpochSampler(Sampler):
+    def __init__(
+        self,
+        *,
+        size: int,
+        sample_count: int,
+        shuffle: bool = False,
+        seed: int = 0,
+        start: Optional[int] = None,
+        step: Optional[int] = None,
+    ):
+        self._size = size
+        self._sample_count = sample_count
+        self._shuffle = shuffle
+        self._seed = seed
+        self._start = get_rank() if start is None else start
+        self._step = get_world_size() if step is None else step
+        self._epoch = 0
+
+    def __iter__(self):
+        count = (self._size + self._sample_count - 1) // self._sample_count
+        tiled_indices = np.tile(np.arange(self._sample_count), count)
+        if self._shuffle:
+            seed = self._seed * self._epoch if self._seed != 0 else self._epoch
+            rng = np.random.default_rng(seed)
+            iterable = rng.choice(tiled_indices, self._size, replace=False)
+        else:
+            iterable = tiled_indices[: self._size]
+
+        yield from itertools.islice(iterable, self._start, None, self._step)
+
+    def __len__(self):
+        return (self._size - self._start + self._step - 1) // self._step
+
+    def set_epoch(self, epoch):
+        self._epoch = epoch
+
+
+def _get_numpy_dtype(size: int) -> Any:
+    return np.int32 if size <= 2**31 else np.int64
+
+
+def _get_torch_dtype(size: int) -> Any:
+    return torch.int32 if size <= 2**31 else torch.int64
+
+
+def _generate_randperm_indices(*, size: int, generator: torch.Generator):
+    """Generate the indices of a random permutation."""
+    dtype = _get_torch_dtype(size)
+    # This is actually matching PyTorch's CPU implementation, see: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/TensorFactories.cpp#L900-L921
+    perm = torch.arange(size, dtype=dtype)
+    for i in range(size):
+        j = torch.randint(i, size, size=(1,), generator=generator).item()
+
+        # Always swap even if no-op
+        value = perm[j].item()
+        perm[j] = perm[i].item()
+        perm[i] = value
+        yield value
+
+
+class InfiniteSampler(Sampler):
+    def __init__(
+        self,
+        *,
+        sample_count: int,
+        shuffle: bool = False,
+        seed: int = 0,
+        start: Optional[int] = None,
+        step: Optional[int] = None,
+        advance: int = 0,
+    ):
+        self._sample_count = sample_count
+        self._seed = seed
+        self._shuffle = shuffle
+        self._start = get_rank() if start is None else start
+        self._step = get_world_size() if step is None else step
+        self._advance = advance
+
+    def __iter__(self):
+        if self._shuffle:
+            iterator = self._shuffled_iterator()
+        else:
+            iterator = self._iterator()
+
+        yield from itertools.islice(iterator, self._advance, None)
+
+    def _iterator(self):
+        assert not self._shuffle
+
+        while True:
+            iterable = range(self._sample_count)
+            yield from itertools.islice(iterable, self._start, None, self._step)
+
+    def _shuffled_iterator(self):
+        assert self._shuffle
+
+        # Instantiate a generator here (rather than in the ctor) to keep the class
+        # picklable (requirement of mp.spawn)
+        generator = torch.Generator().manual_seed(self._seed)
+
+        while True:
+            iterable = _generate_randperm_indices(size=self._sample_count, generator=generator)
+            yield from itertools.islice(iterable, self._start, None, self._step)
+
+
+# The following function is somewhat equivalent to _new_shuffle_tensor_slice below,
+# but avoids a full in-place random permutation generation.
+def _shuffle_tensor_slice(
+    *, tensor: torch.Tensor, start: int = 0, step: int = 1, generator: torch.Generator
+) -> np.ndarray:
+    stop = len(tensor)
+    count = stop // step
+    drop_count = stop - step * count
+    if drop_count:
+        warnings.warn(f"# of dropped samples: {drop_count}", stacklevel=1)
+
+    dtype = _get_numpy_dtype(stop)
+    result = np.empty(count, dtype=dtype)
+
+    for i in range(count):
+        j = torch.randint(0, i + 1, size=(1,), generator=generator).item() if i > 0 else 0
+
+        result[i] = result[j]
+        result[j] = tensor[start + i * step].item()
+
+    return result
+
+
+def _new_shuffle_tensor_slice(
+    *, tensor: torch.Tensor, start: int = 0, step: int = 1, generator: torch.Generator
+) -> np.ndarray:
+    stop = len(tensor)
+    count = stop // step
+    dtype = torch.int64  # Needed for using randperm result as indices
+    count = stop // step
+    drop_count = stop - step * count
+    if drop_count:
+        warnings.warn(f"# of dropped samples: {drop_count}", stacklevel=1)
+    indices = torch.randperm(count, dtype=dtype, generator=generator)
+    return tensor[start::step][indices].numpy()
+
+
+def _make_seed(seed: int, start: int, iter_count: int) -> int:
+    # NOTE: Tried a few variants (including iter_count << 32), this one worked best.
+    return seed + start + (iter_count << 24)
+
+
+class ShardedInfiniteSampler(Sampler):
+    def __init__(
+        self,
+        *,
+        sample_count: int,
+        shuffle: bool = False,
+        seed: int = 0,
+        start: Optional[int] = None,
+        step: Optional[int] = None,
+        advance: int = 0,
+        use_new_shuffle_tensor_slice: bool = False,
+    ):
+        self._sample_count = sample_count
+        self._seed = seed
+        self._shuffle = shuffle
+        self._start = get_rank() if start is None else start
+        self._step = get_world_size() if step is None else step
+        self._advance = advance
+        self._iter_count = 0
+        self._shuffle_tensor_slice_fn = (
+            _new_shuffle_tensor_slice if use_new_shuffle_tensor_slice else _shuffle_tensor_slice
+        )
+
+    def __iter__(self):
+        iter_count = self._advance // self._sample_count
+        if iter_count > 0:
+            self._advance -= iter_count * self._sample_count
+            self._iter_count += iter_count
+
+        if self._shuffle:
+            iterator = self._shuffled_iterator()
+        else:
+            iterator = self._iterator()
+
+        yield from itertools.islice(iterator, self._advance, None)
+
+    def _iterator(self):
+        assert not self._shuffle
+
+        while True:
+            iterable = range(self._sample_count)
+            yield from itertools.islice(iterable, self._start, None, self._step)
+
+    def _shuffled_iterator(self):
+        assert self._shuffle
+
+        # Instantiate a generator here (rather than in the ctor) to be keep the class
+        # picklable (requirement of mp.spawn)
+        generator = torch.Generator()
+
+        # Always shuffle everything first
+        generator.manual_seed(self._seed)
+        dtype = _get_torch_dtype(self._sample_count)
+        perm = torch.randperm(self._sample_count, dtype=dtype, generator=generator)
+
+        while True:
+            # Re-seed on each iteration to allow skipping whole permutations
+            seed = _make_seed(self._seed, self._start, self._iter_count)
+            generator.manual_seed(seed)
+
+            iterable = self._shuffle_tensor_slice_fn(
+                tensor=perm, start=self._start, step=self._step, generator=generator
+            )
+            yield from iterable
+            self._iter_count += 1

dinov3/data/transforms.py

@@ -0,0 +1,165 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import logging
+from typing import Sequence
+
+import torch
+from torchvision.transforms import v2
+
+logger = logging.getLogger("dinov3")
+
+
+def make_interpolation_mode(mode_str: str) -> v2.InterpolationMode:
+    return {mode.value: mode for mode in v2.InterpolationMode}[mode_str]
+
+
+class GaussianBlur(v2.RandomApply):
+    """
+    Apply Gaussian Blur to the PIL image.
+    """
+
+    def __init__(self, *, p: float = 0.5, radius_min: float = 0.1, radius_max: float = 2.0):
+        # NOTE: torchvision is applying 1 - probability to return the original image
+        keep_p = 1 - p
+        transform = v2.GaussianBlur(kernel_size=9, sigma=(radius_min, radius_max))
+        super().__init__(transforms=[transform], p=keep_p)
+
+
+# Use timm's names
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+
+CROP_DEFAULT_SIZE = 224
+RESIZE_DEFAULT_SIZE = int(256 * CROP_DEFAULT_SIZE / 224)
+
+
+def make_normalize_transform(
+    mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
+    std: Sequence[float] = IMAGENET_DEFAULT_STD,
+) -> v2.Normalize:
+    return v2.Normalize(mean=mean, std=std)
+
+
+def make_base_transform(
+    mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
+    std: Sequence[float] = IMAGENET_DEFAULT_STD,
+) -> v2.Normalize:
+    return v2.Compose(
+        [
+            v2.ToDtype(torch.float32, scale=True),
+            make_normalize_transform(mean=mean, std=std),
+        ]
+    )
+
+
+# This roughly matches torchvision's preset for classification training:
+#   https://github.com/pytorch/vision/blob/main/references/classification/presets.py#L6-L44
+def make_classification_train_transform(
+    *,
+    crop_size: int = CROP_DEFAULT_SIZE,
+    interpolation=v2.InterpolationMode.BICUBIC,
+    hflip_prob: float = 0.5,
+    mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
+    std: Sequence[float] = IMAGENET_DEFAULT_STD,
+):
+    transforms_list = [v2.ToImage(), v2.RandomResizedCrop(crop_size, interpolation=interpolation)]
+    if hflip_prob > 0.0:
+        transforms_list.append(v2.RandomHorizontalFlip(hflip_prob))
+    transforms_list.append(make_base_transform(mean, std))
+    transform = v2.Compose(transforms_list)
+    logger.info(f"Built classification train transform\n{transform}")
+    return transform
+
+
+def make_resize_transform(
+    *,
+    resize_size: int,
+    resize_square: bool = False,
+    resize_large_side: bool = False,  # Set the larger side to resize_size instead of the smaller
+    interpolation: v2.InterpolationMode = v2.InterpolationMode.BICUBIC,
+):
+    assert not (resize_square and resize_large_side), "These two options can not be set together"
+    if resize_square:
+        logger.info("resizing image as a square")
+        size = (resize_size, resize_size)
+        transform = v2.Resize(size=size, interpolation=interpolation)
+        return transform
+    elif resize_large_side:
+        logger.info("resizing based on large side")
+        transform = v2.Resize(size=None, max_size=resize_size, interpolation=interpolation)
+        return transform
+    else:
+        transform = v2.Resize(resize_size, interpolation=interpolation)
+        return transform
+
+
+# Derived from make_classification_eval_transform() with more control over resize and crop
+def make_eval_transform(
+    *,
+    resize_size: int = RESIZE_DEFAULT_SIZE,
+    crop_size: int = CROP_DEFAULT_SIZE,
+    resize_square: bool = False,
+    resize_large_side: bool = False,  # Set the larger side to resize_size instead of the smaller
+    interpolation: v2.InterpolationMode = v2.InterpolationMode.BICUBIC,
+    mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
+    std: Sequence[float] = IMAGENET_DEFAULT_STD,
+) -> v2.Compose:
+    transforms_list = [v2.ToImage()]
+    resize_transform = make_resize_transform(
+        resize_size=resize_size,
+        resize_square=resize_square,
+        resize_large_side=resize_large_side,
+        interpolation=interpolation,
+    )
+    transforms_list.append(resize_transform)
+    if crop_size:
+        transforms_list.append(v2.CenterCrop(crop_size))
+    transforms_list.append(make_base_transform(mean, std))
+    transform = v2.Compose(transforms_list)
+    logger.info(f"Built eval transform\n{transform}")
+    return transform
+
+
+# This matches (roughly) torchvision's preset for classification evaluation:
+#   https://github.com/pytorch/vision/blob/main/references/classification/presets.py#L47-L69
+def make_classification_eval_transform(
+    *,
+    resize_size: int = RESIZE_DEFAULT_SIZE,
+    crop_size: int = CROP_DEFAULT_SIZE,
+    interpolation=v2.InterpolationMode.BICUBIC,
+    mean: Sequence[float] = IMAGENET_DEFAULT_MEAN,
+    std: Sequence[float] = IMAGENET_DEFAULT_STD,
+) -> v2.Compose:
+    return make_eval_transform(
+        resize_size=resize_size,
+        crop_size=crop_size,
+        interpolation=interpolation,
+        mean=mean,
+        std=std,
+        resize_square=False,
+        resize_large_side=False,
+    )
+
+
+def voc2007_classification_target_transform(label, n_categories=20):
+    one_hot = torch.zeros(n_categories, dtype=int)
+    for instance in label.instances:
+        one_hot[instance.category_id] = True
+    return one_hot
+
+
+def imaterialist_classification_target_transform(label, n_categories=294):
+    one_hot = torch.zeros(n_categories, dtype=int)
+    one_hot[label.attributes] = True
+    return one_hot
+
+
+def get_target_transform(dataset_str):
+    if "VOC2007" in dataset_str:
+        return voc2007_classification_target_transform
+    elif "IMaterialist" in dataset_str:
+        return imaterialist_classification_target_transform
+    return None

dinov3/distributed/__init__.py

@@ -0,0 +1,24 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+# isort: skip_file
+from .torch_distributed_wrapper import (
+    disable_distributed as disable,
+    enable_distributed as enable,
+    get_default_process_group,
+    get_process_subgroup,
+    get_rank,
+    get_subgroup_rank,
+    get_subgroup_size,
+    get_world_size,
+    is_distributed_enabled as is_enabled,
+    is_main_process,
+    is_subgroup_main_process,
+    new_subgroups,
+    save_in_main_process,
+    TorchDistributedEnvironment,
+)
+
+from .torch_distributed_primitives import gather_all_tensors, reduce_dict

dinov3/distributed/torch_distributed_primitives.py

@@ -0,0 +1,100 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from typing import Any, Dict, List, Optional
+
+import torch
+import torch.distributed as dist
+from torch.nn import functional as F
+
+from .torch_distributed_wrapper import get_default_process_group, get_world_size
+
+
+def reduce_dict(input_dict: Dict[Any, torch.Tensor], average: bool = True) -> Dict[Any, torch.Tensor]:
+    """
+    Reduce the values in the dictionary from all processes so that all processes
+    have the averaged results. Returns a dictionary with the same fields as
+    the input dictionary, after reduction.
+
+    Args:
+        input_dict (dict): all the values will be reduced
+        average (bool): whether to do average or sum
+    """
+    world_size = get_world_size()
+    if world_size <= 1:
+        return input_dict
+    with torch.no_grad():
+        names = []
+        values = []
+        # Sort the keys so that they are consistent across processes
+        for k in sorted(input_dict.keys()):
+            names.append(k)
+            values.append(input_dict[k])
+        stacked_values = torch.stack(values, dim=0)
+        dist.all_reduce(stacked_values)
+        if average:
+            stacked_values /= world_size
+        reduced_dict = {k: v for k, v in zip(names, stacked_values)}
+    return reduced_dict
+
+
+def _simple_gather_all_tensors(result: torch.Tensor, group: Any, world_size: int) -> List[torch.Tensor]:
+    gathered_result = [torch.zeros_like(result) for _ in range(world_size)]
+    dist.all_gather(gathered_result, result, group)
+    return gathered_result
+
+
+def gather_all_tensors(result: torch.Tensor, group: Optional[Any] = None) -> List[torch.Tensor]:
+    """
+    Copied from https://github.com/Lightning-AI/torchmetrics/blob/master/src/torchmetrics/utilities/distributed.py
+    Gather all tensors from several ddp processes onto a list that is broadcasted to all processes.
+
+    Works on tensors that have the same number of dimensions, but where each dimension may differ. In this case
+    tensors are padded, gathered and then trimmed to secure equal workload for all processes.
+
+    Args:
+        result: the value to sync
+        group: the process group to gather results from. Defaults to all processes (world)
+
+    Return:
+        list with size equal to the process group where element i corresponds to result tensor from process i
+    """
+    if group is None:
+        group = get_default_process_group()
+
+    # convert tensors to contiguous format
+    result = result.contiguous()
+
+    world_size = get_world_size()
+    dist.barrier(group=group)
+
+    # if the tensor is scalar, things are easy
+    if result.ndim == 0:
+        return _simple_gather_all_tensors(result, group, world_size)
+
+    # 1. Gather sizes of all tensors
+    local_size = torch.tensor(result.shape, device=result.device)
+    local_sizes = [torch.zeros_like(local_size) for _ in range(world_size)]
+    dist.all_gather(local_sizes, local_size, group=group)
+    max_size = torch.stack(local_sizes).max(dim=0).values
+    all_sizes_equal = all(all(ls == max_size) for ls in local_sizes)
+
+    # 2. If shapes are all the same, then do a simple gather:
+    if all_sizes_equal:
+        return _simple_gather_all_tensors(result, group, world_size)
+
+    # 3. If not, we need to pad each local tensor to maximum size, gather and then truncate
+    pad_dims = []
+    pad_by = (max_size - local_size).detach().cpu()
+    for val in reversed(pad_by):
+        pad_dims.append(0)
+        pad_dims.append(val.item())
+    result_padded = F.pad(result, pad_dims)
+    gathered_result = [torch.zeros_like(result_padded) for _ in range(world_size)]
+    dist.all_gather(gathered_result, result_padded, group)
+    for idx, item_size in enumerate(local_sizes):
+        slice_param = [slice(dim_size) for dim_size in item_size]
+        gathered_result[idx] = gathered_result[idx][slice_param]
+    return gathered_result

dinov3/distributed/torch_distributed_wrapper.py

@@ -0,0 +1,351 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import logging
+import os
+import random
+import socket
+import subprocess
+from datetime import timedelta
+from enum import Enum
+from typing import List, Sequence
+
+import torch
+import torch.distributed as dist
+
+logger = logging.getLogger("dinov3")
+
+_DEFAULT_PROCESS_GROUP = None
+_PROCESS_SUBGROUP = None
+_BUILTIN_PRINT = None
+
+
+def is_distributed_enabled() -> bool:
+    """
+    Returns:
+        True if distributed training is enabled.
+    """
+    return dist.is_available() and dist.is_initialized()
+
+
+def get_rank(group=None) -> int:
+    """
+    Returns:
+        The rank of the current process within the specified process group.
+    """
+    if not is_distributed_enabled():
+        return 0
+    return dist.get_rank(group=group)
+
+
+def get_world_size(group=None) -> int:
+    """
+    Returns:
+        The number of processes in the specified process group.
+    """
+    if not is_distributed_enabled():
+        return 1
+    return dist.get_world_size(group=group)
+
+
+def is_main_process(group=None) -> bool:
+    """
+    Returns:
+        True if the current process is the main one in the specified process group.
+    """
+    return get_rank(group) == 0
+
+
+def save_in_main_process(*args, **kwargs) -> None:
+    """Utility function to save only from the main process."""
+    group = kwargs.pop("group", None)
+    if not is_main_process(group):
+        return
+    torch.save(*args, **kwargs)
+
+
+def _restrict_print_to_main_process() -> None:
+    """This function disables printing when not in the main process."""
+    import builtins as __builtin__
+
+    global _BUILTIN_PRINT
+    _BUILTIN_PRINT = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        if is_main_process() or force:
+            _BUILTIN_PRINT(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def _get_master_port(seed: int = 0) -> int:
+    MIN_MASTER_PORT, MAX_MASTER_PORT = (20_000, 60_000)
+
+    master_port_str = os.environ.get("MASTER_PORT")
+    if master_port_str is None:
+        rng = random.Random(seed)
+        return rng.randint(MIN_MASTER_PORT, MAX_MASTER_PORT)
+
+    return int(master_port_str)
+
+
+def _get_available_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        # A "" host address means INADDR_ANY i.e. binding to all interfaces.
+        # Note this is not compatible with IPv6.
+        s.bind(("", 0))
+        port = s.getsockname()[1]
+        return port
+
+
+def _parse_slurm_node_list(s: str) -> List[str]:
+    return subprocess.check_output(["scontrol", "show", "hostnames", s], text=True).splitlines()
+
+
+class JobType(Enum):
+    TORCHELASTIC = "TorchElastic"
+    SLURM = "Slurm"
+    MANUAL = "manual"
+
+
+class TorchDistributedEnvironment:
+    """
+    Helper class to get (and set) distributed job information from the
+    environment. Identifies and supports (in this order):
+    - TorchElastic,
+    - Slurm,
+    - Manual launch (single-node).
+    """
+
+    def __init__(self):
+        if "TORCHELASTIC_RUN_ID" in os.environ:
+            # TorchElastic job created with torchrun
+            self.job_id = os.environ["TORCHELASTIC_RUN_ID"]
+            self.job_type = JobType.TORCHELASTIC
+
+            self.master_addr = os.environ["MASTER_ADDR"]
+            self.master_port = int(os.environ["MASTER_PORT"])
+            self.rank = int(os.environ["RANK"])
+            self.world_size = int(os.environ["WORLD_SIZE"])
+            self.local_rank = int(os.environ["LOCAL_RANK"])
+            self.local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])
+        elif "SLURM_JOB_ID" in os.environ:
+            # Slurm job created with sbatch, submitit, etc...
+            self.job_id = int(os.environ["SLURM_JOB_ID"])
+            self.job_type = JobType.SLURM
+
+            node_count = int(os.environ["SLURM_JOB_NUM_NODES"])
+            nodes = _parse_slurm_node_list(os.environ["SLURM_JOB_NODELIST"])
+            assert len(nodes) == node_count
+
+            self.master_addr = nodes[0]
+            self.master_port = _get_master_port(seed=self.job_id)
+            self.rank = int(os.environ["SLURM_PROCID"])
+            self.world_size = int(os.environ["SLURM_NTASKS"])
+            self.local_rank = int(os.environ["SLURM_LOCALID"])
+            self.local_world_size = self.world_size // node_count
+        else:
+            # Single node and single job launched manually
+            self.job_id = None
+            self.job_type = JobType.MANUAL
+
+            self.master_addr = "127.0.0.1"
+            self.master_port = _get_available_port()
+            self.rank = 0
+            self.world_size = 1
+            self.local_rank = 0
+            self.local_world_size = 1
+
+        assert self.rank < self.world_size
+        assert self.local_rank < self.local_world_size
+
+    def export(
+        self,
+        *,
+        overwrite: bool,
+        nccl_async_error_handling: bool = False,
+    ) -> "TorchDistributedEnvironment":
+        # See the "Environment variable initialization" section from
+        # https://pytorch.org/docs/stable/distributed.html for the complete list of
+        # environment variables required for the env:// initialization method.
+        env_vars = {
+            "MASTER_ADDR": self.master_addr,
+            "MASTER_PORT": str(self.master_port),
+            "RANK": str(self.rank),
+            "WORLD_SIZE": str(self.world_size),
+            "LOCAL_RANK": str(self.local_rank),
+            "LOCAL_WORLD_SIZE": str(self.local_world_size),
+        }
+        if nccl_async_error_handling:
+            env_vars.update(
+                {
+                    "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1",  # "TORCH_" prefix added in PyTorch 2.2
+                }
+            )
+
+        if not overwrite:
+            for k, v in env_vars.items():
+                # Only check for difference with preset environment variables
+                if k not in os.environ:
+                    continue
+                if os.environ[k] == v:
+                    continue
+                raise RuntimeError(f"Cannot export environment variables as {k} is already set")
+
+        os.environ.update(env_vars)
+        return self
+
+    @property
+    def is_main_process(self) -> bool:
+        return self.rank == 0
+
+    def __str__(self):
+        return (
+            f"{self.job_type.value} job "
+            + (f"({self.job_id}) " if self.job_id else "")
+            + f"using {self.master_addr}:{self.master_port} "  # noqa: E231
+            f"(rank={self.rank}, world size={self.world_size})"
+        )
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}("
+            f"master_addr={self.master_addr},"  # noqa: E231
+            f"master_port={self.master_port},"  # noqa: E231
+            f"rank={self.rank},"  # noqa: E231
+            f"world_size={self.world_size},"  # noqa: E231
+            f"local_rank={self.local_rank},"  # noqa: E231
+            f"local_world_size={self.local_world_size}"
+            ")"
+        )
+
+
+def enable_distributed(
+    *,
+    set_cuda_current_device: bool = True,
+    overwrite: bool = False,
+    nccl_async_error_handling: bool = False,
+    restrict_print_to_main_process: bool = True,
+    timeout: timedelta | None = None,
+):
+    """Enable distributed mode.
+
+    Args:
+        set_cuda_current_device: If True, call torch.cuda.set_device() to set the
+            current PyTorch CUDA device to the one matching the local rank.
+        overwrite: If True, overwrites already set variables. Else fails.
+        nccl_async_error_handling: Enables NCCL asynchronous error handling. As a
+            side effect, this enables timing out PyTorch distributed operations
+            after a default 30 minutes delay).
+        restrict_print_to_main_process: If True, the print function of non-main processes
+            (ie rank>0) is disabled. Use print(..., force=True) to print anyway.
+            If False, nothing is changed and all processes can print as usual.
+        timeout: Timeout for operations executed against the process group.
+            Default value is 10 minutes for NCCL and 30 minutes for other backends.
+    """
+    global _DEFAULT_PROCESS_GROUP
+
+    if _DEFAULT_PROCESS_GROUP is not None:
+        raise RuntimeError("Distributed mode has already been enabled")
+
+    torch_env = TorchDistributedEnvironment()
+    logger.info(f"PyTorch distributed environment: {torch_env}")
+    torch_env.export(
+        overwrite=overwrite,
+        nccl_async_error_handling=nccl_async_error_handling,
+    )
+
+    if set_cuda_current_device:
+        torch.cuda.set_device(torch_env.local_rank)
+
+    dist.init_process_group(backend="nccl", timeout=timeout)
+    dist.barrier()
+
+    if restrict_print_to_main_process:
+        _restrict_print_to_main_process()
+
+    # Finalize setup
+    _DEFAULT_PROCESS_GROUP = torch.distributed.group.WORLD
+
+
+def get_default_process_group():
+    return _DEFAULT_PROCESS_GROUP
+
+
+def disable_distributed() -> None:
+    global _BUILTIN_PRINT
+    if _BUILTIN_PRINT is not None:
+        import builtins as __builtin__
+
+        __builtin__.print = _BUILTIN_PRINT
+
+    global _PROCESS_SUBGROUP
+    # checking here because get_process_subgroup can return _DEFAULT_PROCESS_GROUP
+    if _PROCESS_SUBGROUP is not None:
+        torch.distributed.destroy_process_group(_PROCESS_SUBGROUP)
+        _PROCESS_SUBGROUP = None
+
+    global _DEFAULT_PROCESS_GROUP
+    if _DEFAULT_PROCESS_GROUP is not None:  # not initialized
+        torch.distributed.destroy_process_group(_DEFAULT_PROCESS_GROUP)
+        _DEFAULT_PROCESS_GROUP = None
+
+
+def new_subgroups(all_subgroup_ranks: Sequence[Sequence[int]]):
+    """Create new process subgroups according to the provided specification.
+
+    Args:
+       all_subgroup_ranks: a sequence of rank sequences (first rank, ..., last rank),
+           one for each process subgroup. Example: ((0, 1), (2, 3), (4, 5, 6, 7)).
+
+    Note:
+       This is similar to the (non-documented) new_subgroups_by_enumeration().
+       This should be called once (and not sequentially) to create all subgroups.
+    """
+    all_ranks = tuple(rank for subgroup_ranks in all_subgroup_ranks for rank in subgroup_ranks)
+    rank = get_rank()
+    assert len(all_ranks) == len(set(all_ranks))
+    assert rank in all_ranks
+
+    global _PROCESS_SUBGROUP
+    assert _PROCESS_SUBGROUP is None
+
+    for subgroup_ranks in all_subgroup_ranks:
+        subgroup = torch.distributed.new_group(subgroup_ranks)
+        if rank in subgroup_ranks:
+            _PROCESS_SUBGROUP = subgroup
+
+
+def get_process_subgroup():
+    """
+    Returns:
+        The process subgroup of this rank (or None).
+    """
+    return _PROCESS_SUBGROUP or _DEFAULT_PROCESS_GROUP
+
+
+def get_subgroup_rank() -> int:
+    """
+    Returns:
+        The rank of the current process within its process subgroup.
+    """
+    return get_rank(group=get_process_subgroup())
+
+
+def get_subgroup_size() -> int:
+    """
+    Returns:
+        The number of processes in the process subgroup
+    """
+    return get_world_size(group=get_process_subgroup())
+
+
+def is_subgroup_main_process() -> bool:
+    """
+    Returns:
+        True if the current process is the main one within its process subgroup.
+    """
+    return get_rank(group=get_process_subgroup()) == 0

dinov3/env/__init__.py

@@ -0,0 +1,33 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import contextlib
+import logging
+import os
+import tempfile
+from typing import Optional
+
+import submitit.helpers
+
+logger = logging.getLogger("dinov3")
+
+
+@contextlib.contextmanager
+def clean_env():
+    try:
+        # Hide torch.compile() variables from the launched evals
+        extra_names = ("TRITON_CACHE_DIR", "TORCHINDUCTOR_CACHE_DIR")
+        ctx = submitit.helpers.clean_env(extra_names=extra_names)
+    except TypeError as e:
+        logger.warning("Update submitit to the latest main branch\n%s", e)
+        ctx = submitit.helpers.clean_env()
+    with ctx:
+        yield
+
+
+def set_triton_cache_dir(cache_dir: Optional[str] = None) -> None:
+    if cache_dir is None:
+        cache_dir = tempfile.mkdtemp()
+    os.environ["TRITON_CACHE_DIR"] = cache_dir

dinov3/eval/accumulators.py

@@ -0,0 +1,77 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import logging
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+import torch
+from torch import Tensor
+
+from dinov3.distributed import gather_all_tensors  # Gathers tensors of different sizes
+
+logger = logging.getLogger("dinov3")
+
+
+def _cat_and_gather_tensor_list(tensor_list: List[Tensor]) -> Tensor:
+    local_cat = torch.cat(tensor_list)
+    return torch.cat(gather_all_tensors(local_cat))
+
+
+class Accumulator:
+    def __init__(self) -> None:
+        pass
+
+    def update(self, preds: Tensor, target: Tensor, index: Tensor) -> None:
+        raise NotImplementedError
+
+    def accumulate(self) -> Optional[Dict[str, Tensor]]:
+        raise NotImplementedError
+
+
+class NoOpAccumulator(Accumulator):
+    def __init__(self) -> None:
+        pass
+
+    def update(self, preds: Tensor, target: Tensor, index: Tensor) -> None:
+        pass
+
+    def accumulate(self) -> None:
+        return None
+
+
+class ResultsAccumulator(Accumulator):
+    """
+    Accumulate predictions and targets across processes
+    """
+
+    def __init__(self) -> None:
+        self._local_values: Dict[str, List[Tensor]] = defaultdict(list)
+        self._gathered_values: Dict[str, Tensor] = {}
+        self._gathered = False
+
+    def update(self, preds: Tensor, target: Tensor, index: Tensor) -> None:
+        assert len(preds) == len(target) == len(index)
+        assert not self._gathered, "Tensors have already been gathered in this helper"
+        self._local_values["preds"].append(preds)
+        self._local_values["target"].append(target)
+        self._local_values["index"].append(index)
+        self._gathered = False
+
+    def _gather_tensors(self):
+        for k, tensor_list in self._local_values.items():
+            self._gathered_values[k] = _cat_and_gather_tensor_list(tensor_list)
+        self._gathered = True
+
+    def accumulate(self) -> Dict[str, Tensor]:
+        if not self._gathered:
+            self._gather_tensors()
+        preds, target, index = [self._gathered_values[k] for k in ["preds", "target", "index"]]
+        assert len(preds) == len(target) == len(index) and index.min() == 0
+        preds_ordered = torch.zeros((index.max() + 1, *preds.shape[1:]), dtype=preds.dtype, device=preds.device)
+        preds_ordered[index] = preds
+        target_ordered = torch.zeros((index.max() + 1, *target.shape[1:]), dtype=target.dtype, device=target.device)
+        target_ordered[index] = target
+        return {"preds": preds_ordered, "target": target_ordered}

dinov3/eval/data.py

@@ -0,0 +1,256 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import logging
+from functools import lru_cache
+from typing import Any, Callable, Optional
+
+import numpy as np
+import torch
+from torch.utils.data import Subset
+from torchvision.datasets.vision import StandardTransform
+
+from dinov3.eval.utils import extract_features
+
+logger = logging.getLogger("dinov3")
+
+
+class SubsetEx(Subset):
+    def _get_actual_index(self, index: int) -> int:
+        return self.indices[index]
+
+    def get_target(self, index: int) -> Any:
+        actual_index = self._get_actual_index(index)
+        return self.dataset.get_target(actual_index)
+
+    @property
+    def transforms(self):
+        return self.dataset.transforms
+
+
+def get_target_transform(dataset) -> Optional[Callable]:
+    if hasattr(dataset, "transforms"):
+        if isinstance(dataset.transforms, StandardTransform):
+            return dataset.transforms.target_transform
+        raise ValueError("Dataset has a non-standard .transforms property")
+    if hasattr(dataset, "target_transform"):
+        return dataset.target_transform
+    return None
+
+
+@lru_cache(maxsize=1)
+def get_labels(dataset) -> torch.Tensor:
+    """
+    Get the labels of a classification dataset, as a Tensor, using the `get_targets` method
+    if it is present or loading the labels one by one with `get_target`, if it exists.
+    If the dataset has a target transform, iterate over the whole dataset to get the
+    transformed labels for each element, then stack them as a torch tensor.
+    """
+    logger.info("Getting dataset labels ...")
+    if hasattr(dataset, "get_targets") or hasattr(dataset, "get_target"):
+        if hasattr(dataset, "get_targets"):  # Returns a np.array
+            labels = dataset.get_targets()
+        elif hasattr(dataset, "get_target"):
+            labels = [dataset.get_target(i) for i in range(len(dataset))]
+        target_transform = get_target_transform(dataset)
+        if target_transform is not None:
+            labels = [target_transform(label) for label in labels]
+    else:
+        # Target transform is applied in this case
+        labels = [dataset[i][1] for i in range(len(dataset))]
+    return torch.stack([torch.tensor(label, dtype=int) for label in labels])
+
+
+def get_num_classes(dataset) -> int:
+    """
+    Get the labels of a dataset and compute the number of classes
+    """
+    labels = get_labels(dataset)
+    if len(labels.shape) > 1:
+        return int(labels.shape[1])
+    return int(labels.max() + 1)
+
+
+def create_class_indices_mapping(labels: torch.Tensor) -> dict[int, torch.Tensor]:
+    """
+    Efficiently creates a mapping between the labels and tensors containing
+    the indices of all the dataset elements that share this label.
+    In the case of multiple labels, it is not guaranteed that there
+    will be exactly the specified percentage of labels.
+    """
+    if len(labels.shape) > 1:  # labels are a one-hot encoding
+        assert len(labels.shape) == 2
+        sorted_labels, indices = torch.nonzero(labels.T, as_tuple=True)
+    else:
+        sorted_labels, indices = torch.sort(labels, stable=True)
+    unique_labels, counts = torch.unique_consecutive(sorted_labels, return_counts=True)
+    mapping = dict(zip(unique_labels.tolist(), torch.split(indices, counts.tolist())))
+    return mapping
+
+
+def _shuffle_dataset(dataset: torch.Tensor, seed: int = 0):
+    """
+    Shuffling a dataset by subsetting it with a random permutation of its indices
+    """
+    random_generator = torch.Generator()
+    random_generator.manual_seed(seed)
+    random_indices = torch.randperm(len(dataset), generator=random_generator)
+    return SubsetEx(dataset, random_indices)
+
+
+def _subset_dataset_per_class(
+    class_indices_mapping: dict[int, torch.Tensor],
+    n_or_percent_per_class: int | float,
+    dataset_size: int,
+    seed: int = 0,
+    is_percent: bool = False,
+) -> torch.Tensor:
+    """
+    Helper function to select a percentage of a dataset, equally distributed across classes,
+    or to take the same number of elements from each class of the dataset.
+    Returns a boolean mask tensor being True at indices of selected elements
+    """
+
+    random_generator = torch.Generator()
+    random_generator.manual_seed(seed)
+
+    final_indices_bool = torch.zeros(dataset_size, dtype=bool)
+    for class_indices in class_indices_mapping.values():
+        # Select at least one element
+        n_for_class = max(int(len(class_indices) * n_or_percent_per_class), 1) if is_percent else n_or_percent_per_class
+        assert isinstance(n_for_class, int)
+        filtered_index = torch.randperm(len(class_indices), generator=random_generator)[:n_for_class]
+        final_indices_bool[class_indices[filtered_index]] = True
+    return final_indices_bool
+
+
+def _multilabel_rebalance_subset(
+    class_indices_mapping: dict[int, torch.Tensor],
+    n_or_percent_per_class: int | float,
+    labels: torch.Tensor,
+    indices_bool: torch.Tensor,
+    dataset_size: int,
+    seed: int = 0,
+) -> torch.Tensor:
+    """
+    Helper function to refine a subset of a multi-label dataset (indices_bool)
+    to better match a target percentage of labels.
+    Returns a boolean mask tensor being True at indices of selected elements.
+    """
+
+    # Compute the number of selected labels in indices_bool
+    num_total_labels = labels.sum()
+    num_wanted_labels = int(num_total_labels * n_or_percent_per_class)
+    num_selected_labels = (labels[indices_bool] > 0).sum()
+    logger.info(f" {num_selected_labels} labels instead of {num_wanted_labels}")
+
+    # Compute a new percentage and new set selecting less images, therefore less labels, to match approximatelly the exact percentage of labels selected
+    n_or_percent_per_class = n_or_percent_per_class / (num_selected_labels / num_wanted_labels)
+    final_indices_bool = _subset_dataset_per_class(
+        class_indices_mapping, n_or_percent_per_class, dataset_size, seed, True
+    )
+
+    # Compute the number of labels finally used
+    num_selected_labels = (labels[final_indices_bool] > 0).sum()
+    logger.info(f" {num_selected_labels} labels instead of {num_wanted_labels}")
+
+    return final_indices_bool
+
+
+def split_train_val_datasets(train_dataset, split_percentage: float = 0.1, shuffle_train: bool = True):
+    """
+    Splitting a percent of the train dataset to choose hyperparameters, taking the same percentage for each class.
+    If `shuffle` is False, taking the first elements of each class as the validaton set.
+    """
+    assert 0 < split_percentage < 1
+    logger.info(f"Selecting {int(split_percentage * 100)}% of the train dataset as the validation set")
+    if shuffle_train:
+        logger.info("Shuffling train dataset before splitting in train and validation sets")
+        train_dataset = _shuffle_dataset(train_dataset)
+    train_labels = get_labels(train_dataset)
+    class_indices_mapping = create_class_indices_mapping(train_labels)
+    val_mask = torch.zeros(len(train_labels), dtype=bool)
+    for class_indices in class_indices_mapping.values():
+        # If there is only one element, it goes in the train set
+        n_for_val = max(1, int(split_percentage * len(class_indices))) if len(class_indices) > 1 else 0
+        val_mask[class_indices[:n_for_val]] = True
+
+    val_dataset = SubsetEx(train_dataset, val_mask.nonzero().flatten())
+    train_dataset = SubsetEx(train_dataset, (~val_mask).nonzero().flatten())
+    return train_dataset, val_dataset
+
+
+def create_train_dataset_dict(
+    train_dataset,
+    few_shot_eval: bool = False,
+    few_shot_k_or_percent: float | None = None,
+    few_shot_n_tries: int = 1,
+) -> dict[int, dict[int, Any]]:
+    """
+    Randomly split a dataset for few-shot evaluation, with `few_shot_k_or_percent` being
+    n elements or x% of a class. Produces a dict, which keys are number of random "tries"
+    and values are the dataset subset for this "try".
+
+    Format is {"nth-try": dataset}
+    """
+    if few_shot_eval is False:
+        assert few_shot_k_or_percent is None
+        assert few_shot_n_tries == 1
+        return {0: train_dataset}
+
+    assert few_shot_k_or_percent is not None
+    train_labels = get_labels(train_dataset)
+    class_indices_mapping = create_class_indices_mapping(train_labels)
+    train_dataset_dict: dict[int, Any] = {}
+    is_percent = few_shot_k_or_percent < 1
+    if not is_percent:
+        few_shot_k_or_percent = int(few_shot_k_or_percent)
+
+    for t in range(few_shot_n_tries):
+        t_subset_bool = _subset_dataset_per_class(
+            class_indices_mapping=class_indices_mapping,
+            n_or_percent_per_class=few_shot_k_or_percent,
+            dataset_size=len(train_labels),
+            is_percent=is_percent,
+            seed=t,
+        )
+        if len(train_labels.shape) > 1 and is_percent:
+            t_subset_bool = _multilabel_rebalance_subset(
+                class_indices_mapping=class_indices_mapping,
+                n_or_percent_per_class=few_shot_k_or_percent,
+                dataset_size=len(train_labels),
+                labels=train_labels,
+                indices_bool=t_subset_bool,
+                seed=t,
+            )
+        train_dataset_dict[t] = SubsetEx(train_dataset, t_subset_bool.nonzero().flatten())
+    return train_dataset_dict
+
+
+def extract_features_for_dataset_dict(
+    model, dataset_dict: dict[int, dict[int, Any]], batch_size: int, num_workers: int, gather_on_cpu=False
+) -> dict[int, dict[str, torch.Tensor]]:
+    """
+    Extract features for each subset of dataset in the context of few-shot evaluations
+    """
+    few_shot_data_dict: dict[int, dict[str, torch.Tensor]] = {}
+    for try_n, dataset in dataset_dict.items():
+        features, labels = extract_features(model, dataset, batch_size, num_workers, gather_on_cpu=gather_on_cpu)
+        few_shot_data_dict[try_n] = {"train_features": features, "train_labels": labels}
+    return few_shot_data_dict
+
+
+def pad_multilabel_and_collate(batch, pad_value=-1):
+    """
+    This method pads and collates a batch of (image, (index, target)) tuples, coming from
+    DatasetWithEnumeratedTargets, with targets that are list of potentially varying sizes.
+    The targets are padded to the length of the longest target list in the batch.
+    """
+    maxlen = max(len(targets) for _, (_, targets) in batch)
+    padded_batch = [
+        (image, (index, np.pad(targets, (0, maxlen - len(targets)), constant_values=pad_value)))
+        for image, (index, targets) in batch
+    ]
+    return torch.utils.data.default_collate(padded_batch)

dinov3/eval/dense/depth/models/__init__.py

@@ -0,0 +1,186 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import torch
+from dinov3.eval.dense.depth.utils import cast_to
+
+from .dpt_head import DPTHead
+from .encoder import BackboneLayersSet, DinoVisionTransformerWrapper, PatchSizeAdaptationStrategy
+
+
+class FeaturesToDepth(torch.nn.Module):
+    def __init__(
+        self,
+        min_depth=0.001,
+        max_depth=80,
+        bins_strategy="linear",
+        norm_strategy="linear",
+    ):
+        """
+        Module which converts a feature maps into a depth map
+
+        Args:
+        min_depth (float): minimum depth, used to calibrate the depth range
+        max_depth (float): maximum depth, used to calibrate the depth range
+        bins_strategy (str): Choices are 'linear' or 'log', for Uniform or Scale Invariant distributions for depth bins.
+                             See AdaBins [1] for more details.
+        norm_strategy (str): Choices are 'linear', 'softmax' or 'sigmoid', for the conversion of features to depth logits
+        scale_up (bool): If true, and only if regression by classification is not used, the result is multiplied by max_depth
+
+
+        Example:
+        x = depth_model(input_image)  # N C H W
+        - If pure regression (C == 1), depth is obtained by scaling and/or shifting x
+        - If C > 1, bins are used:
+            Depth is obtained as a weighted sum of depth bins, where weights are predicted logits. (see AdaBins [1] for more details)
+
+        [1] AdaBins: https://github.com/shariqfarooq123/AdaBins
+        """
+        super().__init__()
+        self.min_depth = min_depth
+        self.max_depth = max_depth
+        assert bins_strategy in ["linear", "log"], "Support bins_strategy: linear, log"
+        assert norm_strategy in ["linear", "softmax", "sigmoid"], "Support norm_strategy: linear, softmax, sigmoid"
+
+        self.bins_strategy = bins_strategy
+        self.norm_strategy = norm_strategy
+
+    def forward(self, x):
+        n_bins = x.shape[1]  # N n_bins H W
+        if n_bins > 1:
+            if self.bins_strategy == "linear":
+                bins = torch.linspace(self.min_depth, self.max_depth, n_bins, device=x.device)
+            elif self.bins_strategy == "log":
+                bins = torch.linspace(
+                    torch.log(torch.tensor(self.min_depth)),
+                    torch.log(torch.tensor(self.max_depth)),
+                    n_bins,
+                    device=x.device,
+                )
+                bins = torch.exp(bins)
+
+            # following Adabins, default linear
+            if self.norm_strategy == "linear":
+                logit = torch.relu(x)
+                eps = 0.1
+                logit = logit + eps
+                logit = logit / logit.sum(dim=1, keepdim=True)
+            elif self.norm_strategy == "softmax":
+                logit = torch.softmax(x, dim=1)
+            elif self.norm_strategy == "sigmoid":
+                logit = torch.sigmoid(x)
+                logit = logit / logit.sum(dim=1, keepdim=True)
+
+            output = torch.einsum("ikmn,k->imn", [logit, bins]).unsqueeze(dim=1)
+        else:
+            # standard regression
+            output = torch.relu(x) + self.min_depth
+        return output
+
+
+def make_head(
+    embed_dims: int | list[int],
+    n_output_channels: int,
+    use_batchnorm: bool = False,
+    use_cls_token: bool = False,
+    # upsample: int = 4,
+    head_type: str = "linear",
+    **kwargs,
+) -> torch.nn.Module:
+
+    if isinstance(embed_dims, int):
+        embed_dims = [embed_dims]
+
+    if head_type == "linear":
+        raise NotImplementedError
+    elif head_type == "dpt":
+        decoder = DPTHead(
+            in_channels=embed_dims,
+            n_output_channels=n_output_channels,
+            readout_type="project" if use_cls_token else "ignore",
+            use_batchnorm=use_batchnorm,
+            **kwargs,  # TODO add here post_process_channels, n_hidden_channels
+        )
+    else:
+        raise NotImplementedError("only linear and DPT head supported")
+    return decoder
+
+
+class EncoderDecoder(torch.nn.Module):
+    def __init__(
+        self,
+        encoder: torch.nn.Module,
+        decoder: torch.nn.Module,
+        encoder_dtype=torch.float,
+        decoder_dtype=torch.float,
+    ):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.encoder_dtype = encoder_dtype
+        self.decoder_dtype = decoder_dtype
+        self.is_cuda = torch.cuda.is_available()
+
+    def forward(self, x):
+        with torch.autocast("cuda" if torch.cuda.is_available() else "cpu", enabled=False):
+            x = x.to(self.encoder_dtype)
+            x = self.encoder(x)
+            x = cast_to(x, self.decoder_dtype)
+        return self.decoder(x)
+
+
+def build_depther(
+    backbone: torch.nn.Module,
+    backbone_out_layers: list[int] | BackboneLayersSet,
+    n_output_channels: int,
+    use_backbone_norm: bool = False,
+    use_batchnorm: bool = False,
+    use_cls_token: bool = False,
+    adapt_to_patch_size: PatchSizeAdaptationStrategy = PatchSizeAdaptationStrategy.CENTER_PADDING,
+    head_type: str = "dpt",
+    encoder_dtype: torch.dtype = torch.float,
+    decoder_dtype: torch.dtype = torch.float,
+    # depth args
+    min_depth: float = 0.001,
+    max_depth: float = 10.0,
+    bins_strategy: str = "linear",
+    norm_strategy: str = "linear",
+    **kwargs,
+):
+    encoder = DinoVisionTransformerWrapper(
+        backbone_model=backbone,
+        backbone_out_layers=backbone_out_layers,
+        use_backbone_norm=use_backbone_norm,
+        adapt_to_patch_size=adapt_to_patch_size,
+    )
+    encoder = encoder.to(encoder_dtype)
+    encoder.eval()
+
+    decoder = make_head(
+        encoder.embed_dims,
+        n_output_channels=n_output_channels,
+        use_batchnorm=use_batchnorm,
+        use_cls_token=use_cls_token,
+        head_type=head_type,
+        **kwargs,
+    )
+    decoder.eval()
+
+    features_to_depth = FeaturesToDepth(
+        min_depth=min_depth,
+        max_depth=max_depth,
+        bins_strategy=bins_strategy,
+        norm_strategy=norm_strategy,
+    )
+
+    return torch.nn.Sequential(
+        EncoderDecoder(
+            encoder,
+            decoder,
+            encoder_dtype=encoder_dtype,
+            decoder_dtype=decoder_dtype,
+        ),
+        features_to_depth,
+    )

dinov3/eval/dense/depth/models/dpt_head.py

@@ -0,0 +1,532 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import torch
+from torch import nn
+
+
+def kaiming_init(
+    module: nn.Module,
+    a: float = 0,
+    mode: str = "fan_out",
+    nonlinearity: str = "relu",
+    bias: float = 0,
+    distribution: str = "normal",
+) -> None:
+    assert distribution in ["uniform", "normal"]
+    if hasattr(module, "weight") and module.weight is not None:
+        if distribution == "uniform":
+            nn.init.kaiming_uniform_(module.weight, a=a, mode=mode, nonlinearity=nonlinearity)
+        else:
+            nn.init.kaiming_normal_(module.weight, a=a, mode=mode, nonlinearity=nonlinearity)
+    if hasattr(module, "bias") and module.bias is not None:
+        nn.init.constant_(module.bias, bias)
+
+
+def constant_init(module, val, bias=0):
+    if hasattr(module, "weight") and module.weight is not None:
+        nn.init.constant_(module.weight, val)
+    if hasattr(module, "bias") and module.bias is not None:
+        nn.init.constant_(module.bias, bias)
+
+
+class ConvModule(nn.Module):
+    """A conv block that bundles conv/norm/activation layers.
+    This block simplifies the usage of convolution layers, which are commonly
+    used with a norm layer (e.g., BatchNorm) and activation layer (e.g., ReLU).
+    It is based upon three build methods: `build_conv_layer()`,
+    `build_norm_layer()` and `build_activation_layer()`.
+    Besides, we add some additional features in this module.
+    1. Automatically set `bias` of the conv layer.
+    2. Spectral norm is supported.
+    3. More padding modes are supported. Before PyTorch 1.5, nn.Conv2d only
+    supports zero and circular padding, and we add "reflect" padding mode.
+    Args:
+        in_channels (int): Number of channels in the input feature map.
+            Same as that in ``nn._ConvNd``.
+        out_channels (int): Number of channels produced by the convolution.
+            Same as that in ``nn._ConvNd``.
+        kernel_size (int | tuple[int]): Size of the convolving kernel.
+            Same as that in ``nn._ConvNd``.
+        stride (int | tuple[int]): Stride of the convolution.
+            Same as that in ``nn._ConvNd``.
+        padding (int | tuple[int]): Zero-padding added to both sides of
+            the input. Same as that in ``nn._ConvNd``.
+        dilation (int | tuple[int]): Spacing between kernel elements.
+            Same as that in ``nn._ConvNd``.
+        groups (int): Number of blocked connections from input channels to
+            output channels. Same as that in ``nn._ConvNd``.
+        bias (bool | str): If specified as `auto`, it will be decided by the
+            norm_cfg. Bias will be set as True if `norm_cfg` is None, otherwise
+            False. Default: "auto".
+        conv_cfg (dict): Config dict for convolution layer. Default: None,
+            which means using conv2d.
+        norm_cfg (dict): Config dict for normalization layer. Default: None.
+        act_cfg (dict): Config dict for activation layer.
+            Default: dict(type='ReLU').
+        inplace (bool): Whether to use inplace mode for activation.
+            Default: True.
+        with_spectral_norm (bool): Whether use spectral norm in conv module.
+            Default: False.
+        padding_mode (str): If the `padding_mode` has not been supported by
+            current `Conv2d` in PyTorch, we will use our own padding layer
+            instead. Currently, we support ['zeros', 'circular'] with official
+            implementation and ['reflect'] with our own implementation.
+            Default: 'zeros'.
+        order (tuple[str]): The order of conv/norm/activation layers. It is a
+            sequence of "conv", "norm" and "act". Common examples are
+            ("conv", "norm", "act") and ("act", "conv", "norm").
+            Default: ('conv', 'norm', 'act').
+    """
+
+    _abbr_ = "conv_block"
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups=1,
+        bias="auto",
+        conv_cfg=None,
+        norm_cfg=None,
+        act_cfg=dict(type="ReLU"),
+        inplace=True,
+        with_spectral_norm=False,
+        padding_mode="zeros",
+        order=("conv", "norm", "act"),
+    ):
+        super().__init__()
+        assert conv_cfg is None or isinstance(conv_cfg, dict)
+        assert norm_cfg is None or isinstance(norm_cfg, dict)
+        assert act_cfg is None or isinstance(act_cfg, dict)
+        official_padding_mode = ["zeros", "circular"]
+        self.conv_cfg = conv_cfg
+        self.norm_cfg = norm_cfg
+        self.act_cfg = act_cfg
+        self.inplace = inplace
+        self.with_spectral_norm = with_spectral_norm
+        self.with_explicit_padding = padding_mode not in official_padding_mode
+        self.order = order
+        assert isinstance(self.order, tuple) and len(self.order) == 3
+        assert set(order) == {"conv", "norm", "act"}
+
+        self.with_norm = norm_cfg is not None
+        self.with_activation = act_cfg is not None
+        # if the conv layer is before a norm layer, bias is unnecessary.
+        if bias == "auto":
+            bias = not self.with_norm
+        self.with_bias = bias
+
+        # if self.with_explicit_padding:
+        #     pad_cfg = dict(type=padding_mode)
+        #     self.padding_layer = build_padding_layer(pad_cfg, padding)
+        # to do Camille put back
+
+        # reset padding to 0 for conv module
+        conv_padding = 0 if self.with_explicit_padding else padding
+        # build convolution layer
+        self.conv = nn.Conv2d(  # build_conv_layer(#conv_cfg,
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=conv_padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+        )
+        # export the attributes of self.conv to a higher level for convenience
+        self.in_channels = self.conv.in_channels
+        self.out_channels = self.conv.out_channels
+        self.kernel_size = self.conv.kernel_size
+        self.stride = self.conv.stride
+        self.padding = padding
+        self.dilation = self.conv.dilation
+        self.transposed = self.conv.transposed
+        self.output_padding = self.conv.output_padding
+        self.groups = self.conv.groups
+
+        if self.with_spectral_norm:
+            self.conv = nn.utils.spectral_norm(self.conv)
+
+        # # build normalization layers
+        if self.with_norm:
+            # norm layer is after conv layer
+            if order.index("norm") > order.index("conv"):
+                norm_channels = out_channels
+            else:
+                norm_channels = in_channels
+            # self.norm_name, norm = build_norm_layer(
+            #     norm_cfg, norm_channels)  # type: ignore
+            self.add_module("bn", torch.nn.SyncBatchNorm(norm_channels))
+            # if self.with_bias:
+            #     if isinstance(norm, (_BatchNorm, _InstanceNorm)):
+            #         warnings.warn(
+            #             'Unnecessary conv bias before batch/instance norm')
+            self.norm_name = "bn"
+        else:
+            self.norm_name = None  # type: ignore
+
+        # build activation layer
+        if self.with_activation:
+            act_cfg_ = act_cfg.copy()  # type: ignore
+            # nn.Tanh has no 'inplace' argument
+            if act_cfg_["type"] not in ["Tanh", "PReLU", "Sigmoid", "HSigmoid", "Swish", "GELU"]:
+                act_cfg_.setdefault("inplace", inplace)
+            self.activate = nn.ReLU()  # build_activation_layer(act_cfg_)
+
+        # Use msra init by default
+        self.init_weights()
+
+    @property
+    def norm(self):
+        if self.norm_name:
+            return getattr(self, self.norm_name)
+        else:
+            return None
+
+    def init_weights(self):
+        # 1. It is mainly for customized conv layers with their own
+        #    initialization manners by calling their own ``init_weights()``,
+        #    and we do not want ConvModule to override the initialization.
+        # 2. For customized conv layers without their own initialization
+        #    manners (that is, they don't have their own ``init_weights()``)
+        #    and PyTorch's conv layers, they will be initialized by
+        #    this method with default ``kaiming_init``.
+        # Note: For PyTorch's conv layers, they will be overwritten by our
+        #    initialization implementation using default ``kaiming_init``.
+        if not hasattr(self.conv, "init_weights"):
+            if self.with_activation and self.act_cfg["type"] == "LeakyReLU":
+                nonlinearity = "leaky_relu"
+                a = self.act_cfg.get("negative_slope", 0.01)
+            else:
+                nonlinearity = "relu"
+                a = 0
+            kaiming_init(self.conv, a=a, nonlinearity=nonlinearity)
+        if self.with_norm:
+            constant_init(self.norm_name, 1, bias=0)
+
+    def forward(self, x: torch.Tensor, activate: bool = True, norm: bool = True, debug: bool = False) -> torch.Tensor:
+        for layer in self.order:
+            if debug:
+                breakpoint()
+            if layer == "conv":
+                if self.with_explicit_padding:
+                    x = self.padding_layer(x)
+                x = self.conv(x)
+            elif layer == "norm" and norm and self.with_norm:
+                x = self.norm(x)
+            elif layer == "act" and activate and self.with_activation:
+                x = self.activate(x)
+        return x
+
+
+class Interpolate(nn.Module):
+    def __init__(self, scale_factor, mode, align_corners=False):
+        super(Interpolate, self).__init__()
+        self.interp = nn.functional.interpolate
+        self.scale_factor = scale_factor
+        self.mode = mode
+        self.align_corners = align_corners
+
+    def forward(self, x):
+        x = self.interp(x, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners)
+        return x
+
+
+class UpConvHead(nn.Module):
+    """
+    A 3 layer Convolutional head with intermediate upsampling
+
+    Args:
+    - features (int): number of input channels
+    - n_output_channels (int, default=256): number of output channels
+    - n_hidden_channels (int, default=32): number of channels in hidden layer
+
+    The operations are
+    [
+        Conv3x3(features, features // 2),
+        2x-Upsampling,
+        Conv3x3(features // 2, hidden_channels),
+        ReLU,
+        Conv1x1(hidden_channels, n_output_channels),
+    ]
+    """
+
+    def __init__(self, features, n_output_channels, n_hidden_channels=32):
+        super(UpConvHead, self).__init__()
+        self.n_output_channels = n_output_channels
+        self.head = nn.Sequential(
+            nn.Conv2d(features, features // 2, kernel_size=3, stride=1, padding=1),
+            Interpolate(scale_factor=2, mode="bilinear", align_corners=True),
+            nn.Conv2d(features // 2, n_hidden_channels, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(n_hidden_channels, n_output_channels, kernel_size=1, stride=1, padding=0),
+        )
+
+    def forward(self, x):
+        x = self.head(x)
+        return x
+
+
+class ReassembleBlocks(nn.Module):
+    """ViTPostProcessBlock, process cls_token in ViT backbone output and
+    rearrange the feature vector to feature map.
+    Args:
+        in_channels (List): ViT feature channels.
+            Default: [1024, 1024, 1024, 1024].
+        out_channels (List): output channels of each stage.
+            Default: [128, 256, 512, 1024].
+        readout_type (str): Type of readout operation. Default: 'ignore'.
+        init_cfg (dict, optional): Initialization config dict. Default: None.
+    """
+
+    def __init__(
+        self,
+        in_channels=[1024, 1024, 1024, 1024],
+        out_channels=[128, 256, 512, 1024],
+        readout_type="project",
+        use_batchnorm=False,
+    ):
+        super(ReassembleBlocks, self).__init__()
+
+        assert readout_type in ["ignore", "add", "project"]
+        self.readout_type = readout_type
+
+        self.projects = nn.ModuleList(
+            [
+                ConvModule(
+                    in_channels=in_channels[channel_index],
+                    out_channels=out_channel,
+                    kernel_size=1,
+                    act_cfg=None,
+                )
+                for channel_index, out_channel in enumerate(out_channels)
+            ]
+        )
+
+        self.resize_layers = nn.ModuleList(
+            [
+                nn.ConvTranspose2d(
+                    in_channels=out_channels[0], out_channels=out_channels[0], kernel_size=4, stride=4, padding=0
+                ),
+                nn.ConvTranspose2d(
+                    in_channels=out_channels[1], out_channels=out_channels[1], kernel_size=2, stride=2, padding=0
+                ),
+                nn.Identity(),
+                nn.Conv2d(
+                    in_channels=out_channels[3], out_channels=out_channels[3], kernel_size=3, stride=2, padding=1
+                ),
+            ]
+        )
+        if self.readout_type == "project":
+            self.readout_projects = nn.ModuleList()
+            for i in range(len(self.projects)):
+                self.readout_projects.append(nn.Sequential(nn.Linear(2 * in_channels[i], in_channels[i]), nn.GELU()))
+
+        self.batchnorm_layers = nn.ModuleList(
+            [nn.SyncBatchNorm(channel) if use_batchnorm else nn.Identity(channel) for channel in in_channels]
+        )
+
+    def forward(self, inputs):
+        assert isinstance(inputs, list)
+        out = []
+        for i, x in enumerate(inputs):
+            assert len(x) == 2
+            x, cls_token = x[0], x[1]
+            feature_shape = x.shape
+            if self.readout_type == "project":
+                x = x.flatten(2).permute((0, 2, 1))
+                readout = cls_token.unsqueeze(1).expand_as(x)
+                x = self.readout_projects[i](torch.cat((x, readout), -1))
+                x = x.permute(0, 2, 1).reshape(feature_shape)
+            elif self.readout_type == "add":
+                x = x.flatten(2) + cls_token.unsqueeze(-1)
+                x = x.reshape(feature_shape)
+            else:
+                pass
+            x = self.batchnorm_layers[i](x)
+            x = self.projects[i](x)
+            x = self.resize_layers[i](x)
+            out.append(x)
+        return out
+
+
+class PreActResidualConvUnit(nn.Module):
+    """ResidualConvUnit, pre-activate residual unit.
+    Args:
+        in_channels (int): number of channels in the input feature map.
+        act_cfg (dict): dictionary to construct and config activation layer.
+        norm_cfg (dict): dictionary to construct and config norm layer.
+        stride (int): stride of the first block. Default: 1
+        dilation (int): dilation rate for convs layers. Default: 1.
+        init_cfg (dict, optional): Initialization config dict. Default: None.
+    """
+
+    def __init__(self, in_channels, act_cfg, norm_cfg, stride=1, dilation=1, init_cfg=None):
+        super(PreActResidualConvUnit, self).__init__()  # init_cfg)
+        self.conv1 = ConvModule(
+            in_channels,
+            in_channels,
+            3,
+            stride=stride,
+            padding=dilation,
+            dilation=dilation,
+            norm_cfg=norm_cfg,
+            act_cfg=act_cfg,
+            bias=False,
+            order=("act", "conv", "norm"),
+        )
+        self.conv2 = ConvModule(
+            in_channels,
+            in_channels,
+            3,
+            padding=1,
+            norm_cfg=norm_cfg,
+            act_cfg=act_cfg,
+            bias=False,
+            order=("act", "conv", "norm"),
+        )
+
+    def forward(self, inputs):
+        inputs_ = inputs.clone()
+        x = self.conv1(inputs)
+        x = self.conv2(x)
+        return x + inputs_
+
+
+class FeatureFusionBlock(nn.Module):
+    """FeatureFusionBlock, merge feature map from different stages.
+    Args:
+        in_channels (int): Input channels.
+        act_cfg (dict): The activation config for ResidualConvUnit.
+        norm_cfg (dict): Config dict for normalization layer.
+        expand (bool): Whether expand the channels in post process block.
+            Default: False.
+        align_corners (bool): align_corner setting for bilinear upsample.
+            Default: True.
+        init_cfg (dict, optional): Initialization config dict. Default: None.
+    """
+
+    def __init__(self, in_channels, act_cfg, norm_cfg, expand=False, align_corners=True, init_cfg=None):
+        super(FeatureFusionBlock, self).__init__()  # init_cfg)
+        self.in_channels = in_channels
+        self.expand = expand
+        self.align_corners = align_corners
+        self.out_channels = in_channels
+        if self.expand:
+            self.out_channels = in_channels // 2
+        self.project = ConvModule(self.in_channels, self.out_channels, kernel_size=1, act_cfg=None, bias=True)
+        self.res_conv_unit1 = PreActResidualConvUnit(in_channels=self.in_channels, act_cfg=act_cfg, norm_cfg=norm_cfg)
+        self.res_conv_unit2 = PreActResidualConvUnit(in_channels=self.in_channels, act_cfg=act_cfg, norm_cfg=norm_cfg)
+
+    def forward(self, *inputs):
+        x = inputs[0]
+
+        if len(inputs) == 2:
+            if x.shape != inputs[1].shape:
+                res = torch.nn.functional.interpolate(
+                    inputs[1],
+                    size=(x.shape[2], x.shape[3]),
+                    mode="bilinear",
+                    align_corners=False,
+                )
+            else:
+                res = inputs[1]
+            x = x + self.res_conv_unit1(res)
+        x = self.res_conv_unit2(x)  # ok
+
+        x = torch.nn.functional.interpolate(x, scale_factor=2, mode="bilinear", align_corners=self.align_corners)
+        #  ok
+
+        x = self.project(x)  # ok
+        return x
+
+
+class DPTHead(nn.Module):
+    """Vision Transformers for Dense Prediction.
+    This head is implemented of `DPT <https://arxiv.org/abs/2103.13413>`_.
+    Args:
+        in_channels (List): The input dimensions of the ViT backbone.
+            Default: [1024, 1024, 1024, 1024].
+        channels (int): Channels after modules, before the task-specific module
+            (`conv_depth`). Default: 256.
+        post_process_channels (List): Out channels of post process conv
+            layers. Default: [96, 192, 384, 768].
+        readout_type (str): Type of readout operation. Default: 'ignore'.
+        expand_channels (bool): Whether expand the channels in post process
+            block. Default: False.
+    """
+
+    def __init__(
+        self,
+        in_channels=(1024, 1024, 1024, 1024),
+        channels=256,
+        post_process_channels=[128, 256, 512, 1024],
+        readout_type="project",
+        expand_channels=False,
+        n_output_channels=256,
+        use_batchnorm=False,  # TODO
+        **kwargs,
+    ):
+        super(DPTHead, self).__init__(**kwargs)
+        self.channels = channels
+        self.n_output_channels = n_output_channels
+        self.in_channels = in_channels
+        self.expand_channels = expand_channels
+        self.norm_cfg = None  # TODO CHECK THIS
+        self.reassemble_blocks = ReassembleBlocks(
+            in_channels=in_channels,
+            out_channels=post_process_channels,
+            readout_type=readout_type,
+            use_batchnorm=use_batchnorm,
+        )
+
+        self.post_process_channels = [
+            channel * (2**i) if expand_channels else channel for i, channel in enumerate(post_process_channels)
+        ]
+        self.convs = nn.ModuleList()
+        for channel in self.post_process_channels:
+            self.convs.append(ConvModule(channel, self.channels, kernel_size=3, padding=1, act_cfg=None, bias=False))
+        self.fusion_blocks = nn.ModuleList()
+        self.act_cfg = {"type": "ReLU"}
+        for _ in range(len(self.convs)):
+            self.fusion_blocks.append(FeatureFusionBlock(self.channels, self.act_cfg, self.norm_cfg))
+        self.fusion_blocks[0].res_conv_unit1 = None
+        self.project = ConvModule(self.channels, self.channels, kernel_size=3, padding=1, norm_cfg=self.norm_cfg)
+        self.num_fusion_blocks = len(self.fusion_blocks)
+        self.num_reassemble_blocks = len(self.reassemble_blocks.resize_layers)
+        self.num_post_process_channels = len(self.post_process_channels)
+        assert self.num_fusion_blocks == self.num_reassemble_blocks
+        assert self.num_reassemble_blocks == self.num_post_process_channels
+        self.conv_depth = UpConvHead(self.channels, self.n_output_channels)
+
+    def forward_features(self, inputs):
+        assert (
+            len(inputs) == self.num_reassemble_blocks
+        ), f"Expected {self.num_reassemble_blocks} inputs, got {len(inputs)}."
+        x = [inp for inp in inputs]
+
+        x = self.reassemble_blocks(x)
+        x = [self.convs[i](feature) for i, feature in enumerate(x)]
+        out = self.fusion_blocks[0](x[-1])
+
+        for i in range(1, len(self.fusion_blocks)):
+            out = self.fusion_blocks[i](out, x[-(i + 1)])
+
+        out = self.project(out)
+        return out
+
+    def forward(self, inputs):
+        out = self.forward_features(inputs)
+        return self.conv_depth(out)
+
+    def predict(self, inputs, rescale_to=(512, 512)):
+        out = self.forward_features(inputs)
+        return self.conv_depth.predict(out, rescale_to=rescale_to)

dinov3/eval/dense/depth/models/embed.py

@@ -0,0 +1,52 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import itertools
+import math
+
+import torch
+
+
+class CenterPadding(torch.nn.Module):
+    def __init__(self, multiple: int):
+        super().__init__()
+        self.multiple = multiple
+
+    def _get_pad(self, size):
+        new_size = math.ceil(size / self.multiple) * self.multiple
+        pad_size = new_size - size
+        pad_size_left = pad_size // 2
+        pad_size_right = pad_size - pad_size_left
+        return pad_size_left, pad_size_right
+
+    # @torch.inference_mode()
+    def forward(self, x):
+        # expected shapes are ... x H x W, usually B x C x H x W
+        pads = list(itertools.chain.from_iterable(self._get_pad(m) for m in x.shape[:-3:-1]))
+        output = torch.nn.functional.pad(x, pads)
+        return output
+
+    def __extra_repr__(self) -> str:
+        return f"multiple={self.multiple}"
+
+
+class StretchToMultiple(torch.nn.Module):
+    def __init__(self, multiple: int):
+        super().__init__()
+        self.multiple = multiple
+
+    def forward(self, x):
+        # expected shapes are ... x H x W, usually B x C x H x W
+        *shape, C, H, W = x.shape
+        new_H = math.ceil(H / self.multiple) * self.multiple
+        new_W = math.ceil(W / self.multiple) * self.multiple
+        if new_H != H or new_W != W:
+            x = x.reshape(-1, C, H, W)
+            x = torch.nn.functional.interpolate(x, size=(new_H, new_W), mode="bilinear")
+            x = x.reshape(*shape, C, new_H, new_W)
+        return x
+
+    def __extra_repr__(self) -> str:
+        return f"multiple={self.multiple}"

dinov3/eval/dense/depth/models/encoder.py

@@ -0,0 +1,124 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import logging
+from enum import Enum
+from typing import Sequence, Union
+
+from dinov3.eval.dense.depth.models.embed import CenterPadding, StretchToMultiple
+from torch import Tensor, nn
+
+logger = logging.getLogger("dinov3")
+
+
+class BackboneLayersSet(Enum):
+    # Set of intermediate layers to take from the backbone
+    LAST = "LAST"  # extracting only the last layer
+    FOUR_LAST = "FOUR_LAST"  # extracting the last 4 layers
+    FOUR_EVEN_INTERVALS = "FOUR_EVEN_INTERVALS"  # extracting outputs every 1/4 of the total number of blocks
+
+
+def _get_backbone_out_indices(
+    model: nn.Module,
+    backbone_out_layers: Union[list[int], BackboneLayersSet] = BackboneLayersSet.FOUR_EVEN_INTERVALS,
+):
+    """
+    Get indices for output layers of the ViT backbone. For now there are 3 options available:
+    BackboneLayersSet.LAST : only extract the last layer, used in segmentation tasks with a bn head.
+    BackboneLayersSet.FOUR_LAST : extract the last 4 layers, used in segmentation (multiscale setting)
+    BackboneLayersSet.FOUR_EVEN_INTERVALS : extract outputs every 1/4 of the total number of blocks
+    Reference outputs in 'FOUR_EVEN_INTERVALS' mode :
+    ViT/S (12 blocks): [2, 5, 8, 11]
+    ViT/B (12 blocks): [2, 5, 8, 11]
+    ViT/L (24 blocks): [5, 11, 17, 23] (correct), [4, 11, 17, 23] (incorrect)
+    ViT/g (40 blocks): [9, 19, 29, 39]
+    """
+    n_blocks = getattr(model, "n_blocks", 1)
+    if isinstance(backbone_out_layers, list):
+        out_indices = backbone_out_layers
+
+    if backbone_out_layers == BackboneLayersSet.LAST:
+        out_indices = [n_blocks - 1]
+    elif backbone_out_layers == BackboneLayersSet.FOUR_LAST:
+        out_indices = [i for i in range(n_blocks - 4, n_blocks)]
+    elif backbone_out_layers == BackboneLayersSet.FOUR_EVEN_INTERVALS:
+        # XXX: Force (incorrect) out indices for backward-compatibility (ViT/L only)
+        if n_blocks == 24:
+            out_indices = [4, 11, 17, 23]
+        else:
+            out_indices = [i * (n_blocks // 4) - 1 for i in range(1, 5)]
+    assert all([out_index < n_blocks for out_index in out_indices])
+    return out_indices
+
+
+class PatchSizeAdaptationStrategy(Enum):
+    CENTER_PADDING = "center_padding"
+    STRETCH = "stretch"
+    NO_ADAPTATION = "never"
+
+
+class DinoVisionTransformerWrapper(nn.Module):
+    """Vision Transformer."""
+
+    def __init__(
+        self,
+        backbone_model: nn.Module,
+        backbone_out_layers: Union[str, list[int]],
+        use_backbone_norm: bool = False,
+        adapt_to_patch_size: PatchSizeAdaptationStrategy = PatchSizeAdaptationStrategy.CENTER_PADDING,
+    ):
+        super().__init__()
+
+        self.final_norm = use_backbone_norm
+        self.backbone = backbone_model
+        self.backbone_out_indices = _get_backbone_out_indices(
+            self.backbone,
+            backbone_out_layers=(
+                backbone_out_layers if isinstance(backbone_out_layers, list) else BackboneLayersSet(backbone_out_layers)
+            ),
+        )
+
+        # If the backbone does not define embed_dims, use [embed_dim] * n_blocks
+        try:
+            embed_dims = self.backbone.embed_dims
+        except AttributeError:
+            embed_dim = self.backbone.embed_dim
+            n_blocks = self.backbone.n_blocks
+            logger.warning(f"Backbone does not define embed_dims, using {[embed_dim] * n_blocks=} instead")
+            embed_dims = [embed_dim] * n_blocks
+        self.embed_dims: Sequence[int] = [embed_dims[idx] for idx in self.backbone_out_indices]
+
+        # How to adapt input images to the patch size of the model?
+        try:
+            input_pad_size = self.backbone.input_pad_size
+        except AttributeError:
+            patch_size = self.backbone.patch_size
+            logger.warning(f"Backbone does not define input_pad_size, using {patch_size=} instead")
+            input_pad_size = patch_size
+        if adapt_to_patch_size is PatchSizeAdaptationStrategy.CENTER_PADDING:
+            self.patch_size_adapter = CenterPadding(input_pad_size)
+        elif adapt_to_patch_size is PatchSizeAdaptationStrategy.STRETCH:
+            self.patch_size_adapter = StretchToMultiple(input_pad_size)
+        elif adapt_to_patch_size is PatchSizeAdaptationStrategy.NO_ADAPTATION:
+            self.patch_size_adapter = nn.Identity()
+        else:
+            raise ValueError(f"Unknown value {adapt_to_patch_size=}")
+
+        # Freeze backbone
+        self.backbone.requires_grad_(False)
+
+    def forward(
+        self,
+        x: Tensor,  # [B, rgb, H, W]
+    ) -> list[tuple[Tensor, Tensor]]:
+        x = self.patch_size_adapter(x)
+        outputs = self.backbone.get_intermediate_layers(
+            x,
+            n=self.backbone_out_indices,
+            reshape=True,
+            return_class_token=True,
+            norm=self.final_norm,
+        )  # List of (patch feats [B, C, h, w], class token [B, C])
+        return outputs

dinov3/eval/dense/depth/utils.py

@@ -0,0 +1,19 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+import torch
+
+def cast_to(o, dtype: torch.dtype):
+    if isinstance(o, torch.Tensor):
+        return o.to(dtype)
+    if isinstance(o, tuple):
+        return tuple(cast_to(u, dtype) for u in o)
+    if isinstance(o, list):
+        return list(cast_to(u, dtype) for u in o)
+    if isinstance(o, dict):
+        return {k: cast_to(u, dtype) for k, u in o.items()}
+    if isinstance(o, (bool, int, float, str, type(None))):
+        return o
+    raise NotImplementedError(f"Unsupported type: {type(o)}")

dinov3/eval/detection/config.py

@@ -0,0 +1,65 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+from dataclasses import dataclass
+
+from .models.position_encoding import PositionEncoding
+
+
+@dataclass(kw_only=True)
+class DetectionHeadConfig:
+    num_classes: int = 91  # 91 classes in COCO
+    # Deformable DETR tricks
+    with_box_refine: bool = True
+    two_stage: bool = True
+    # DINO DETR tricks
+    mixed_selection: bool = True
+    look_forward_twice: bool = True  # was default False
+    # Hybrid Matching tricks
+    k_one2many: int = 6  # was 5
+    lambda_one2many: float = 1.0
+    num_queries_one2one: int = 300  # number of query slots for one_to_one matching
+    num_queries_one2many: int = 1500  # was 0, number of query slots for one_to_many matching
+    """
+    Absolute coordinates & box regression reparameterization.
+    If true, we use absolute coordindates & reparameterization for bounding boxes.
+    """
+    reparam: bool = True
+    topk: int = 100
+
+    # * Backbone
+    # type of positional embedding to use on top of the image features
+    position_embedding: PositionEncoding = PositionEncoding.SINE
+    num_feature_levels: int = 1  # number of feature levels
+
+    # * Transformer
+    dec_layers: int = 6  # number of decoding layers in the transformer
+    dim_feedforward: int = 2048  # intermediate size of the feedforward layers in the transformer blocks
+    hidden_dim: int = 256  # size of the embeddings (dimension of the transformer)
+    dropout: float = 0.0  # dropout applied in the transformer, was 0.1
+    nheads: int = 8  # number of attention heads inside the transformer's attentions
+    norm_type: str = "pre_norm"
+
+    # Loss
+    aux_loss: bool = True  # auxiliary decoding losses (loss at each layer)
+
+    # * dev: proposals
+    proposal_feature_levels: int = 4  # was 1
+    proposal_min_size: int = 50
+    # * dev decoder: global decoder
+    decoder_type: str = "global_rpe_decomp"  # was deform
+    decoder_use_checkpoint: bool = False
+    decoder_rpe_hidden_dim: int = 512
+    decoder_rpe_type: str = "linear"
+
+    # Custom
+    add_transformer_encoder: bool = True
+    num_encoder_layers: int = 6
+    layers_to_use: list[int] | None = None
+    blocks_to_train: list[int] | None = None
+    n_windows_sqrt: int = 0
+    proposal_in_stride: int | None = None
+    proposal_tgt_strides: list[int] | None = None
+    backbone_use_layernorm: bool = False  # whether to use layernorm on each layer of the backbone's features

dinov3/eval/detection/models/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+

dinov3/eval/detection/models/backbone.py

@@ -0,0 +1,117 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+# ------------------------------------------------------------------------
+# Plain-DETR
+# Copyright (c) 2023 Xi'an Jiaotong University & Microsoft Research Asia.
+# Licensed under The MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modified from DETR (https://github.com/facebookresearch/detr)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# ------------------------------------------------------------------------
+
+"""
+Backbone modules.
+"""
+import logging
+from typing import List, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ..util.misc import NestedTensor
+from .position_encoding import build_position_encoding
+from .utils import LayerNorm2D
+from .windows import WindowsWrapper
+
+logger = logging.getLogger("dinov3")
+
+
+class DINOBackbone(nn.Module):
+    def __init__(
+        self,
+        backbone_model: nn.Module,
+        train_backbone: bool,
+        blocks_to_train: Optional[List[str]] = None,
+        layers_to_use: Union[int, List] = 1,
+        use_layernorm: bool = True,
+    ):
+        super().__init__()
+        self.backbone = backbone_model
+        self.blocks_to_train = blocks_to_train
+        self.patch_size = self.backbone.patch_size
+        self.use_layernorm = use_layernorm
+
+        for _, (name, parameter) in enumerate(self.backbone.named_parameters()):
+            train_condition = any(f".{b}." in name for b in self.blocks_to_train) if self.blocks_to_train else True
+            if (not train_backbone) or "mask_token" in name or (not train_condition):
+                parameter.requires_grad_(False)
+
+        self.strides = [self.backbone.patch_size]
+
+        # get embed_dim for each intermediate output
+        n_all_layers = self.backbone.n_blocks
+        blocks_to_take = (
+            range(n_all_layers - layers_to_use, n_all_layers) if isinstance(layers_to_use, int) else layers_to_use
+        )
+
+        # if models do not define embed_dims, repeat embed_dim n_blocks times
+        embed_dims = getattr(self.backbone, "embed_dims", [self.backbone.embed_dim] * self.backbone.n_blocks)
+        embed_dims = [embed_dims[i] for i in range(n_all_layers) if i in blocks_to_take]
+
+        if self.use_layernorm:
+            self.layer_norms = nn.ModuleList([LayerNorm2D(embed_dim) for embed_dim in embed_dims])
+
+        self.num_channels = [sum(embed_dims)]
+        self.layers_to_use = layers_to_use
+
+    def forward(self, tensor_list: NestedTensor):
+        xs = self.backbone.get_intermediate_layers(tensor_list.tensors, n=self.layers_to_use, reshape=True)
+        if self.use_layernorm:
+            xs = [ln(x).contiguous() for ln, x in zip(self.layer_norms, xs)]
+
+        xs = [torch.cat(xs, axis=1)]
+
+        out: list[NestedTensor] = []
+        for x in xs:
+            m = tensor_list.mask
+            assert m is not None
+            mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0]
+            out.append(NestedTensor(x, mask))
+        return out
+
+
+class BackboneWithPositionEncoding(nn.Sequential):
+    def __init__(self, backbone, position_embedding):
+        super().__init__(backbone, position_embedding)
+        self.strides = backbone.strides
+        self.num_channels = backbone.num_channels
+
+    def forward(self, tensor_list: NestedTensor):
+        out: List[NestedTensor] = list(self[0](tensor_list))
+        pos = [self[1][idx](x).to(x.tensors.dtype) for idx, x in enumerate(out)]
+        return out, pos
+
+
+def build_backbone(backbone_model, args):
+    position_embedding = build_position_encoding(args)
+    train_backbone = False
+    backbone = DINOBackbone(
+        backbone_model, train_backbone, args.blocks_to_train, args.layers_to_use, args.backbone_use_layernorm
+    )
+    if args.n_windows_sqrt > 0:
+        logger.info(f"Wrapping with {args.n_windows_sqrt} x {args.n_windows_sqrt} windows")
+        backbone = WindowsWrapper(
+            backbone, n_windows_w=args.n_windows_sqrt, n_windows_h=args.n_windows_sqrt, patch_size=backbone.patch_size
+        )
+    else:
+        logger.info("Not wrapping with windows")
+
+    return BackboneWithPositionEncoding(backbone, position_embedding)

dinov3/eval/detection/models/detr.py

@@ -0,0 +1,462 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+# ------------------------------------------------------------------------
+# Plain-DETR
+# Copyright (c) 2023 Xi'an Jiaotong University & Microsoft Research Asia.
+# Licensed under The MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modified from DETR (https://github.com/facebookresearch/detr)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# ------------------------------------------------------------------------
+
+"""
+Deformable DETR model and criterion classes.
+"""
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ..util import box_ops
+from ..util.misc import NestedTensor, _get_clones, inverse_sigmoid, nested_tensor_from_tensor_list
+from .backbone import build_backbone
+from .transformer import build_transformer
+
+
+class PlainDETR(nn.Module):
+    """This is the Deformable DETR module that performs object detection"""
+
+    def __init__(
+        self,
+        backbone,
+        transformer,
+        num_classes,
+        num_feature_levels,
+        aux_loss=True,
+        with_box_refine=False,
+        two_stage=False,
+        num_queries_one2one=300,
+        num_queries_one2many=0,
+        mixed_selection=False,
+    ):
+        """Initializes the model.
+        Parameters:
+            backbone: torch module of the backbone to be used. See backbone.py
+            transformer: torch module of the transformer architecture. See transformer.py
+            num_classes: number of object classes
+            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
+            with_box_refine: iterative bounding box refinement
+            two_stage: two-stage Deformable DETR
+            num_queries_one2one: number of object queries for one-to-one matching part
+            num_queries_one2many: number of object queries for one-to-many matching part
+            mixed_selection: a trick for Deformable DETR two stage
+
+        """
+        super().__init__()
+        num_queries = num_queries_one2one + num_queries_one2many
+        self.num_queries = num_queries
+        self.transformer = transformer
+        hidden_dim = transformer.d_model
+        self.class_embed = nn.Linear(hidden_dim, num_classes)
+        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
+        self.num_feature_levels = num_feature_levels
+        if not two_stage:
+            self.query_embed = nn.Embedding(num_queries, hidden_dim * 2)
+        elif mixed_selection:
+            self.query_embed = nn.Embedding(num_queries, hidden_dim)
+        self.input_proj = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.Conv2d(backbone.num_channels[0], hidden_dim, kernel_size=1),
+                    nn.GroupNorm(32, hidden_dim),
+                )
+            ]
+        )
+        self.backbone = backbone
+        self.aux_loss = aux_loss
+        self.with_box_refine = with_box_refine
+        self.two_stage = two_stage
+
+        prior_prob = 0.01
+        bias_value = -math.log((1 - prior_prob) / prior_prob)
+        self.class_embed.bias.data = torch.ones(num_classes) * bias_value
+        nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0)
+        nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0)
+        for proj in self.input_proj:
+            nn.init.xavier_uniform_(proj[0].weight, gain=1)
+            nn.init.constant_(proj[0].bias, 0)
+
+        # if two-stage, the last class_embed and bbox_embed is for region proposal generation
+        num_pred = (transformer.decoder.num_layers + 1) if two_stage else transformer.decoder.num_layers
+        if with_box_refine:
+            self.class_embed = _get_clones(self.class_embed, num_pred)
+            self.bbox_embed = _get_clones(self.bbox_embed, num_pred)
+            nn.init.constant_(self.bbox_embed[0].layers[-1].bias.data[2:], -2.0)
+            # hack implementation for iterative bounding box refinement
+            self.transformer.decoder.bbox_embed = self.bbox_embed
+        else:
+            nn.init.constant_(self.bbox_embed.layers[-1].bias.data[2:], -2.0)
+            self.class_embed = nn.ModuleList([self.class_embed for _ in range(num_pred)])
+            self.bbox_embed = nn.ModuleList([self.bbox_embed for _ in range(num_pred)])
+            self.transformer.decoder.bbox_embed = None
+        if two_stage:
+            # hack implementation for two-stage
+            self.transformer.decoder.class_embed = self.class_embed
+            for box_embed in self.bbox_embed:
+                nn.init.constant_(box_embed.layers[-1].bias.data[2:], 0.0)
+
+        self.num_queries_one2one = num_queries_one2one
+        self.mixed_selection = mixed_selection
+
+    def forward(self, samples: NestedTensor):
+        """The forward expects a NestedTensor, which consists of:
+           - samples.tensor: batched images, of shape [batch_size x 3 x H x W]
+           - samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels
+
+        It returns a dict with the following elements:
+           - "pred_logits": the classification logits (including no-object) for all queries.
+                            Shape= [batch_size x num_queries x (num_classes + 1)]
+           - "pred_boxes": The normalized boxes coordinates for all queries, represented as
+                           (center_x, center_y, height, width). These values are normalized in [0, 1],
+                           relative to the size of each individual image (disregarding possible padding).
+                           See PostProcess for information on how to retrieve the unnormalized bounding box.
+           - "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list of
+                            dictionnaries containing the two above keys for each decoder layer.
+        """
+        if not isinstance(samples, NestedTensor):
+            samples = nested_tensor_from_tensor_list(samples)
+        features, pos = self.backbone(samples)
+
+        srcs = []
+        masks = []
+        for layer, feat in enumerate(features):
+            src, mask = feat.decompose()
+            srcs.append(self.input_proj[layer](src))
+            masks.append(mask)
+            assert mask is not None
+
+        query_embeds = None
+        if not self.two_stage or self.mixed_selection:
+            query_embeds = self.query_embed.weight[0 : self.num_queries, :]
+
+        # make attn mask
+        """ attention mask to prevent information leakage
+        """
+        self_attn_mask = torch.zeros(
+            [
+                self.num_queries,
+                self.num_queries,
+            ],
+            dtype=bool,
+            device=src.device,
+        )
+        self_attn_mask[
+            self.num_queries_one2one :,
+            0 : self.num_queries_one2one,
+        ] = True
+        self_attn_mask[
+            0 : self.num_queries_one2one,
+            self.num_queries_one2one :,
+        ] = True
+
+        (
+            hs,
+            init_reference,
+            inter_references,
+            enc_outputs_class,
+            enc_outputs_coord_unact,
+            enc_outputs_delta,
+            output_proposals,
+            max_shape,
+        ) = self.transformer(srcs, masks, pos, query_embeds, self_attn_mask)
+
+        outputs_classes_one2one = []
+        outputs_coords_one2one = []
+        outputs_classes_one2many = []
+        outputs_coords_one2many = []
+        for lvl in range(hs.shape[0]):
+            if lvl == 0:
+                reference = init_reference
+            else:
+                reference = inter_references[lvl - 1]
+            reference = inverse_sigmoid(reference)
+            outputs_class = self.class_embed[lvl](hs[lvl])
+            tmp = self.bbox_embed[lvl](hs[lvl])
+            if reference.shape[-1] == 4:
+                tmp += reference
+            else:
+                assert reference.shape[-1] == 2
+                tmp[..., :2] += reference
+            outputs_coord = tmp.sigmoid()
+
+            outputs_classes_one2one.append(outputs_class[:, 0 : self.num_queries_one2one])
+            outputs_classes_one2many.append(outputs_class[:, self.num_queries_one2one :])
+
+            outputs_coords_one2one.append(outputs_coord[:, 0 : self.num_queries_one2one])
+            outputs_coords_one2many.append(outputs_coord[:, self.num_queries_one2one :])
+
+        outputs_classes_one2one = torch.stack(outputs_classes_one2one)
+        outputs_coords_one2one = torch.stack(outputs_coords_one2one)
+
+        outputs_classes_one2many = torch.stack(outputs_classes_one2many)
+        outputs_coords_one2many = torch.stack(outputs_coords_one2many)
+
+        out = {
+            "pred_logits": outputs_classes_one2one[-1],
+            "pred_boxes": outputs_coords_one2one[-1],
+            "pred_logits_one2many": outputs_classes_one2many[-1],
+            "pred_boxes_one2many": outputs_coords_one2many[-1],
+        }
+        if self.aux_loss:
+            out["aux_outputs"] = self._set_aux_loss(outputs_classes_one2one, outputs_coords_one2one)
+            out["aux_outputs_one2many"] = self._set_aux_loss(outputs_classes_one2many, outputs_coords_one2many)
+
+        if self.two_stage:
+            enc_outputs_coord = enc_outputs_coord_unact.sigmoid()
+            out["enc_outputs"] = {
+                "pred_logits": enc_outputs_class,
+                "pred_boxes": enc_outputs_coord,
+            }
+        return out
+
+    @torch.jit.unused
+    def _set_aux_loss(self, outputs_class, outputs_coord):
+        # this is a workaround to make torchscript happy, as torchscript
+        # doesn't support dictionary with non-homogeneous values, such
+        # as a dict having both a Tensor and a list.
+        return [{"pred_logits": a, "pred_boxes": b} for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]
+
+
+class PlainDETRReParam(PlainDETR):
+    def forward(self, samples: NestedTensor):
+        """The forward expects a NestedTensor, which consists of:
+           - samples.tensor: batched images, of shape [batch_size x 3 x H x W]
+           - samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels
+
+        It returns a dict with the following elements:
+           - "pred_logits": the classification logits (including no-object) for all queries.
+                            Shape= [batch_size x num_queries x (num_classes + 1)]
+           - "pred_boxes": The normalized boxes coordinates for all queries, represented as
+                           (center_x, center_y, height, width). These values are normalized in [0, 1],
+                           relative to the size of each individual image (disregarding possible padding).
+                           See PostProcess for information on how to retrieve the unnormalized bounding box.
+           - "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list of
+                            dictionnaries containing the two above keys for each decoder layer.
+        """
+        if not isinstance(samples, NestedTensor):
+            samples = nested_tensor_from_tensor_list(samples)
+        features, pos = self.backbone(samples)
+
+        srcs = []
+        masks = []
+        for layer, feat in enumerate(features):
+            src, mask = feat.decompose()
+            srcs.append(self.input_proj[layer](src))
+            masks.append(mask)
+            assert mask is not None
+
+        query_embeds = None
+        if not self.two_stage or self.mixed_selection:
+            query_embeds = self.query_embed.weight[0 : self.num_queries, :]
+
+        # make attn mask
+        """ attention mask to prevent information leakage
+        """
+        self_attn_mask = torch.zeros(
+            [
+                self.num_queries,
+                self.num_queries,
+            ],
+            dtype=bool,
+            device=src.device,
+        )
+        self_attn_mask[
+            self.num_queries_one2one :,
+            0 : self.num_queries_one2one,
+        ] = True
+        self_attn_mask[
+            0 : self.num_queries_one2one,
+            self.num_queries_one2one :,
+        ] = True
+
+        (
+            hs,
+            init_reference,
+            inter_references,
+            enc_outputs_class,
+            enc_outputs_coord_unact,
+            enc_outputs_delta,
+            output_proposals,
+            max_shape,
+        ) = self.transformer(srcs, masks, pos, query_embeds, self_attn_mask)
+
+        outputs_classes_one2one = []
+        outputs_coords_one2one = []
+        outputs_classes_one2many = []
+        outputs_coords_one2many = []
+
+        outputs_coords_old_one2one = []
+        outputs_deltas_one2one = []
+        outputs_coords_old_one2many = []
+        outputs_deltas_one2many = []
+
+        for lvl in range(hs.shape[0]):
+            if lvl == 0:
+                reference = init_reference
+            else:
+                reference = inter_references[lvl - 1]
+            outputs_class = self.class_embed[lvl](hs[lvl])
+            tmp = self.bbox_embed[lvl](hs[lvl])
+            if reference.shape[-1] == 4:
+                outputs_coord = box_ops.box_xyxy_to_cxcywh(box_ops.delta2bbox(reference, tmp, max_shape))
+            else:
+                raise NotImplementedError
+
+            outputs_classes_one2one.append(outputs_class[:, 0 : self.num_queries_one2one])
+            outputs_classes_one2many.append(outputs_class[:, self.num_queries_one2one :])
+
+            outputs_coords_one2one.append(outputs_coord[:, 0 : self.num_queries_one2one])
+            outputs_coords_one2many.append(outputs_coord[:, self.num_queries_one2one :])
+
+            outputs_coords_old_one2one.append(reference[:, : self.num_queries_one2one])
+            outputs_coords_old_one2many.append(reference[:, self.num_queries_one2one :])
+            outputs_deltas_one2one.append(tmp[:, : self.num_queries_one2one])
+            outputs_deltas_one2many.append(tmp[:, self.num_queries_one2one :])
+
+        outputs_classes_one2one = torch.stack(outputs_classes_one2one)
+        outputs_coords_one2one = torch.stack(outputs_coords_one2one)
+
+        outputs_classes_one2many = torch.stack(outputs_classes_one2many)
+        outputs_coords_one2many = torch.stack(outputs_coords_one2many)
+
+        out = {
+            "pred_logits": outputs_classes_one2one[-1],
+            "pred_boxes": outputs_coords_one2one[-1],
+            "pred_logits_one2many": outputs_classes_one2many[-1],
+            "pred_boxes_one2many": outputs_coords_one2many[-1],
+            "pred_boxes_old": outputs_coords_old_one2one[-1],
+            "pred_deltas": outputs_deltas_one2one[-1],
+            "pred_boxes_old_one2many": outputs_coords_old_one2many[-1],
+            "pred_deltas_one2many": outputs_deltas_one2many[-1],
+        }
+
+        if self.aux_loss:
+            out["aux_outputs"] = self._set_aux_loss(
+                outputs_classes_one2one, outputs_coords_one2one, outputs_coords_old_one2one, outputs_deltas_one2one
+            )
+            out["aux_outputs_one2many"] = self._set_aux_loss(
+                outputs_classes_one2many, outputs_coords_one2many, outputs_coords_old_one2many, outputs_deltas_one2many
+            )
+
+        if self.two_stage:
+            out["enc_outputs"] = {
+                "pred_logits": enc_outputs_class,
+                "pred_boxes": enc_outputs_coord_unact,
+                "pred_boxes_old": output_proposals,
+                "pred_deltas": enc_outputs_delta,
+            }
+        return out
+
+    @torch.jit.unused
+    def _set_aux_loss(self, outputs_class, outputs_coord, outputs_coord_old, outputs_deltas):
+        # this is a workaround to make torchscript happy, as torchscript
+        # doesn't support dictionary with non-homogeneous values, such
+        # as a dict having both a Tensor and a list.
+        return [
+            {
+                "pred_logits": a,
+                "pred_boxes": b,
+                "pred_boxes_old": c,
+                "pred_deltas": d,
+            }
+            for a, b, c, d in zip(outputs_class[:-1], outputs_coord[:-1], outputs_coord_old[:-1], outputs_deltas[:-1])
+        ]
+
+
+class PostProcess(nn.Module):
+    """This module converts the model's output into the format expected by the coco api"""
+
+    def __init__(self, topk=100, reparam=False):
+        super().__init__()
+        self.topk = topk
+        self.reparam = reparam
+
+    @torch.no_grad()
+    def forward(self, outputs, target_sizes, original_target_sizes=None):
+        """Perform the computation
+        Parameters:
+            outputs: raw outputs of the model
+            target_sizes: tensor of dimension [batch_size x 2] containing the size of each images of the batch
+                          For evaluation, this must be the original image size (before any data augmentation)
+                          For visualization, this should be the image size after data augment, but before padding
+        """
+        out_logits, out_bbox = outputs["pred_logits"], outputs["pred_boxes"]
+
+        assert len(out_logits) == len(target_sizes)
+        assert target_sizes.shape[1] == 2
+        assert not self.reparam or original_target_sizes.shape[1] == 2
+
+        prob = out_logits.sigmoid()
+        topk_values, topk_indexes = torch.topk(prob.view(out_logits.shape[0], -1), self.topk, dim=1)
+        scores = topk_values
+        topk_boxes = topk_indexes // out_logits.shape[2]
+        labels = topk_indexes % out_logits.shape[2]
+        boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
+        boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))
+
+        # and from relative [0, 1] to absolute [0, height] coordinates
+        img_h, img_w = target_sizes.unbind(1)
+        if self.reparam:
+            img_h, img_w = img_h[:, None, None], img_w[:, None, None]  # [BS, 1, 1]
+            boxes[..., 0::2].clamp_(min=torch.zeros_like(img_w), max=img_w)
+            boxes[..., 1::2].clamp_(min=torch.zeros_like(img_h), max=img_h)
+            scale_h, scale_w = (original_target_sizes / target_sizes).unbind(1)
+            scale_fct = torch.stack([scale_w, scale_h, scale_w, scale_h], dim=1)
+        else:
+            scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
+        boxes = boxes * scale_fct[:, None, :]
+
+        results = [{"scores": s, "labels": l, "boxes": b} for s, l, b in zip(scores, labels, boxes)]
+
+        return results
+
+
+class MLP(nn.Module):
+    """Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+
+
+def build_model(backbone_model, args):
+    backbone = build_backbone(backbone_model, args)
+    transformer = build_transformer(args)
+    model_class = PlainDETR if (not args.reparam) else PlainDETRReParam
+    return model_class(
+        backbone,
+        transformer,
+        num_classes=args.num_classes,
+        num_feature_levels=args.num_feature_levels,
+        aux_loss=args.aux_loss,
+        with_box_refine=args.with_box_refine,
+        two_stage=args.two_stage,
+        num_queries_one2one=args.num_queries_one2one,
+        num_queries_one2many=args.num_queries_one2many,
+        mixed_selection=args.mixed_selection,
+    )

dinov3/eval/detection/models/global_ape_decoder.py

@@ -0,0 +1,321 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+
+# ------------------------------------------------------------------------
+# Plain-DETR
+# Copyright (c) 2023 Xi'an Jiaotong University & Microsoft Research Asia.
+# Licensed under The MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------
+# -*- coding: utf-8 -*-
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
+
+from ..util.misc import _get_activation_fn, _get_clones, inverse_sigmoid
+
+
+class GlobalCrossAttention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        qkv_bias=True,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.q = nn.Linear(dim, dim, bias=qkv_bias)
+        self.k = nn.Linear(dim, dim, bias=qkv_bias)
+        self.v = nn.Linear(dim, dim, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(
+        self,
+        query,
+        k_input_flatten,
+        v_input_flatten,
+        input_padding_mask=None,
+    ):
+        B_, N, C = k_input_flatten.shape
+        k = self.k(k_input_flatten).reshape(B_, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+        v = self.v(v_input_flatten).reshape(B_, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+        B_, N, C = query.shape
+        q = self.q(query).reshape(B_, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+
+        attn_mask = None
+        if input_padding_mask is not None:
+            attn_mask = input_padding_mask[:, None, None] * -100
+        attn_mask = attn_mask.contiguous()  # to enable efficient attention
+
+        x = torch.nn.functional.scaled_dot_product_attention(
+            query=q,
+            key=k,
+            value=v,
+            attn_mask=attn_mask,
+            dropout_p=self.attn_drop.p if self.training else 0,
+            scale=self.scale,
+        )
+        x = x.transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class GlobalDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        d_model=256,
+        d_ffn=1024,
+        dropout=0.1,
+        activation="relu",
+        n_heads=8,
+        norm_type="post_norm",
+    ):
+        super().__init__()
+
+        self.norm_type = norm_type
+
+        # global cross attention
+        self.cross_attn = GlobalCrossAttention(d_model, n_heads)
+        self.dropout1 = nn.Dropout(dropout)
+        self.norm1 = nn.LayerNorm(d_model)
+
+        # self attention
+        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.norm2 = nn.LayerNorm(d_model)
+
+        # ffn
+        self.linear1 = nn.Linear(d_model, d_ffn)
+        self.activation = _get_activation_fn(activation)
+        self.dropout3 = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(d_ffn, d_model)
+        self.dropout4 = nn.Dropout(dropout)
+        self.norm3 = nn.LayerNorm(d_model)
+
+    @staticmethod
+    def with_pos_embed(tensor, pos):
+        return tensor if pos is None else tensor + pos
+
+    def forward_pre(
+        self,
+        tgt,
+        query_pos,
+        src,
+        src_pos_embed,
+        src_padding_mask=None,
+        self_attn_mask=None,
+    ):
+        # self attention
+        tgt2 = self.norm2(tgt)
+        q = k = self.with_pos_embed(tgt2, query_pos)
+        tgt2 = self.self_attn(
+            q.transpose(0, 1), k.transpose(0, 1), tgt2.transpose(0, 1), attn_mask=self_attn_mask, need_weights=False
+        )[0].transpose(0, 1)
+        tgt = tgt + self.dropout2(tgt2)
+
+        # global cross attention
+        tgt2 = self.norm1(tgt)
+        tgt2 = self.cross_attn(
+            self.with_pos_embed(tgt2, query_pos),
+            self.with_pos_embed(src, src_pos_embed),
+            src,
+            src_padding_mask,
+        )
+        tgt = tgt + self.dropout1(tgt2)
+
+        # ffn
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout3(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout4(tgt2)
+
+        return tgt
+
+    def forward_post(
+        self,
+        tgt,
+        query_pos,
+        src,
+        src_pos_embed,
+        src_padding_mask=None,
+        self_attn_mask=None,
+    ):
+        # self attention
+        q = k = self.with_pos_embed(tgt, query_pos)
+        tgt2 = self.self_attn(
+            q.transpose(0, 1), k.transpose(0, 1), tgt.transpose(0, 1), attn_mask=self_attn_mask, need_weights=False
+        )[0].transpose(0, 1)
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+
+        # cross attention
+        tgt2 = self.cross_attn(
+            self.with_pos_embed(tgt, query_pos),
+            self.with_pos_embed(src, src_pos_embed),
+            src,
+            src_padding_mask,
+        )
+        tgt = tgt + self.dropout1(tgt2)
+        tgt = self.norm1(tgt)
+
+        # ffn
+        tgt2 = self.linear2(self.dropout3(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout4(tgt2)
+        tgt = self.norm3(tgt)
+
+        return tgt
+
+    def forward(
+        self,
+        tgt,
+        query_pos,
+        src,
+        src_pos_embed,
+        src_padding_mask=None,
+        self_attn_mask=None,
+    ):
+        if self.norm_type == "pre_norm":
+            return self.forward_pre(tgt, query_pos, src, src_pos_embed, src_padding_mask, self_attn_mask)
+        if self.norm_type == "post_norm":
+            return self.forward_post(tgt, query_pos, src, src_pos_embed, src_padding_mask, self_attn_mask)
+
+
+class GlobalDecoder(nn.Module):
+    def __init__(
+        self,
+        decoder_layer,
+        num_layers,
+        return_intermediate=False,
+        look_forward_twice=False,
+        use_checkpoint=False,
+        d_model=256,
+        norm_type="post_norm",
+    ):
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.return_intermediate = return_intermediate
+        self.look_forward_twice = look_forward_twice
+        self.use_checkpoint = use_checkpoint
+        # hack implementation for iterative bounding box refinement and two-stage Deformable DETR
+        self.bbox_embed = None
+        self.class_embed = None
+
+        self.norm_type = norm_type
+        if self.norm_type == "pre_norm":
+            self.final_layer_norm = nn.LayerNorm(d_model)
+        else:
+            self.final_layer_norm = None
+
+    def _reset_parameters(self):
+        # stolen from Swin Transformer
+        def _init_weights(m):
+            if isinstance(m, nn.Linear):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if isinstance(m, nn.Linear) and m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.LayerNorm):
+                nn.init.constant_(m.bias, 0)
+                nn.init.constant_(m.weight, 1.0)
+
+        self.apply(_init_weights)
+
+    def forward(
+        self,
+        tgt,
+        reference_points,
+        src,
+        src_pos_embed,
+        src_spatial_shapes,
+        src_level_start_index,
+        src_valid_ratios,
+        query_pos=None,
+        src_padding_mask=None,
+        self_attn_mask=None,
+        max_shape=None,
+    ):
+        output = tgt
+
+        intermediate = []
+        intermediate_reference_points = []
+        for lid, layer in enumerate(self.layers):
+            if self.use_checkpoint:
+                output = checkpoint.checkpoint(
+                    layer,
+                    output,
+                    query_pos,
+                    src,
+                    src_pos_embed,
+                    src_padding_mask,
+                    self_attn_mask,
+                )
+            else:
+                output = layer(
+                    output,
+                    query_pos,
+                    src,
+                    src_pos_embed,
+                    src_padding_mask,
+                    self_attn_mask,
+                )
+
+            if self.final_layer_norm is not None:
+                output_after_norm = self.final_layer_norm(output)
+            else:
+                output_after_norm = output
+
+            # hack implementation for iterative bounding box refinement
+            if self.bbox_embed is not None:
+                tmp = self.bbox_embed[lid](output_after_norm)
+                if reference_points.shape[-1] == 4:
+                    new_reference_points = tmp + inverse_sigmoid(reference_points)
+                    new_reference_points = new_reference_points.sigmoid()
+                else:
+                    assert reference_points.shape[-1] == 2
+                    new_reference_points = tmp
+                    new_reference_points[..., :2] = tmp[..., :2] + inverse_sigmoid(reference_points)
+                    new_reference_points = new_reference_points.sigmoid()
+                reference_points = new_reference_points.detach()
+
+            if self.return_intermediate:
+                intermediate.append(output_after_norm)
+                intermediate_reference_points.append(
+                    new_reference_points if self.look_forward_twice else reference_points
+                )
+
+        if self.return_intermediate:
+            return torch.stack(intermediate), torch.stack(intermediate_reference_points)
+
+        return output_after_norm, reference_points
+
+
+def build_global_ape_decoder(args):
+    decoder_layer = GlobalDecoderLayer(
+        d_model=args.hidden_dim,
+        d_ffn=args.dim_feedforward,
+        dropout=args.dropout,
+        activation="relu",
+        n_heads=args.nheads,
+        norm_type=args.norm_type,
+    )
+    decoder = GlobalDecoder(
+        decoder_layer,
+        num_layers=args.dec_layers,
+        return_intermediate=True,
+        look_forward_twice=args.look_forward_twice,
+        use_checkpoint=args.decoder_use_checkpoint,
+        d_model=args.hidden_dim,
+        norm_type=args.norm_type,
+    )
+    return decoder