diff --git a/DeepForestVisionV2.pth b/DeepForestVisionV2.pth
new file mode 100644
index 0000000000000000000000000000000000000000..d71a3d6acc5a2f2fa4a9c76c5b0f5b6dcd451fda
--- /dev/null
+++ b/DeepForestVisionV2.pth
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21f12d6a6610efd8c69f04302b6dc7cdec3c8e75a8306e70348469af60c2c44e
+size 343226287
diff --git a/LICENSE.md b/LICENSE.md
new file mode 100644
index 0000000000000000000000000000000000000000..f531b1e6b5ab2318957bbf8ad1bda9f800a23e17
--- /dev/null
+++ b/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. Grant of Rights. 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. Redistribution and Use.
+
+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.
diff --git a/dinov3/.DS_Store b/dinov3/.DS_Store
new file mode 100644
index 0000000000000000000000000000000000000000..dfd7d801df38c0bfeb23a011f05792a9439c4f09
Binary files /dev/null and b/dinov3/.DS_Store differ
diff --git a/dinov3/__init__.py b/dinov3/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b70f615df55246ea150f0c76e02b0b8b98d16b2b
--- /dev/null
+++ b/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"
diff --git a/dinov3/checkpointer/__init__.py b/dinov3/checkpointer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..120712ccd2ec68848a027ca2a54b662f4e109205
--- /dev/null
+++ b/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,
+)
diff --git a/dinov3/checkpointer/checkpointer.py b/dinov3/checkpointer/checkpointer.py
new file mode 100644
index 0000000000000000000000000000000000000000..7a984110c68be859f5612be5d290bf27d5a77a55
--- /dev/null
+++ b/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}")
diff --git a/dinov3/configs/__init__.py b/dinov3/configs/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..3719887dfb7c33ad67c2546e33e2f1c9f426760c
--- /dev/null
+++ b/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,
+)
diff --git a/dinov3/configs/config.py b/dinov3/configs/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..6aceaa86f9b952cc554cb6ae80ed16a242d67d0e
--- /dev/null
+++ b/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()
diff --git a/dinov3/configs/ssl_default_config.yaml b/dinov3/configs/ssl_default_config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..87ae3a0a33887c8f501d4d94f724f06942bab2a5
--- /dev/null
+++ b/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
diff --git a/dinov3/configs/train/dinov3_vit7b16_gram_anchor.yaml b/dinov3/configs/train/dinov3_vit7b16_gram_anchor.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..05cf7dfef6fda0199e10b45a761f14201a3db6b5
--- /dev/null
+++ b/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
diff --git a/dinov3/configs/train/dinov3_vit7b16_high_res_adapt.yaml b/dinov3/configs/train/dinov3_vit7b16_high_res_adapt.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..313a104a7cbf757cc0288f10d8d22eb960bfe644
--- /dev/null
+++ b/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
diff --git a/dinov3/configs/train/dinov3_vit7b16_pretrain.yaml b/dinov3/configs/train/dinov3_vit7b16_pretrain.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..7b93be46ecedb582befb9d4fcb7ed185410f2517
--- /dev/null
+++ b/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
diff --git a/dinov3/configs/train/dinov3_vitl16_lvd1689m_distilled.yaml b/dinov3/configs/train/dinov3_vitl16_lvd1689m_distilled.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..5d81e605593401de00ea6ac68630714a2c707fcc
--- /dev/null
+++ b/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:
+ output_dir: