项目阅读路线

DINOv3 项目阅读路线

目的：自顶向下梳理 DINOv3 仓库代码，方便对照论文理解实现。 阅读顺序由"主干 → 子模块 → 训练总装 → 损失"展开，可与论文章节一一对应。

1. 项目目录骨架

dinov3/
├── models/             ← 网络主体
│   ├── vision_transformer.py     主干 ViT (DinoVisionTransformer)
│   └── convnext.py               蒸馏时的 ConvNext 教师/学生
├── layers/             ← 网络子模块
│   ├── patch_embed.py            Conv2d patch embed
│   ├── rope_position_encoding.py 2D RoPE
│   ├── attention.py              SelfAttention (含 RoPE 应用)
│   ├── block.py                  SelfAttentionBlock (含 drop_path)
│   ├── ffn_layers.py             Mlp / SwiGLUFFN
│   ├── layer_scale.py / rms_norm.py
│   ├── dino_head.py              DINO/iBOT 共用的 3-layer MLP + 归一化 last layer
│   └── fp8_linear.py / sparse_linear.py
├── loss/               ← 4 个 SSL 损失
│   ├── dino_clstoken_loss.py     DINO 损失 + Sinkhorn-Knopp
│   ├── ibot_patch_loss.py        iBOT 掩码 patch 损失
│   ├── koleo_loss.py             KoLeo 正则
│   └── gram_loss.py              Gram 矩阵一致性损失
├── train/              ← 训练主程序
│   ├── ssl_meta_arch.py          SSLMetaArch：组装 student/teacher/gram_teacher + 损失
│   ├── train.py                  do_train 主循环
│   ├── param_groups.py           分层学习率 / layer-wise decay
│   └── cosine_lr_scheduler.py
├── configs/
│   ├── ssl_default_config.yaml   默认 SSL 配置
│   └── train/*.yaml              不同规模实验配置
├── hub/backbones.py    ← 推理入口 dinov3_vitb16(...) 工厂
├── data/ / eval/       ← 数据加载与评估
└── fsdp/ checkpointer/ distributed/  工程化相关

2. 阅读顺序（与论文对齐）

Step 0｜入口与配置

• dinov3/models/init.py 的 build_model()：用 arch 字符串查表构造 ViT/ConvNext；init_fp8 决定是否把 Linear 转成 FP8。build_model_from_cfg() 在 SSL 训练时用。
• dinov3/configs/ssl_default_config.yaml：看 DINO/iBOT/gram/crops/optim 的默认超参，论文里大部分数字都能在这里对上。
• dinov3/configs/config.py 的 apply_scaling_rules_to_cfg()：理解 sqrt_wrt_1024 学习率缩放。

Step 1｜主干网络：ViT

dinov3/models/vision_transformer.py 是核心。

DinoVisionTransformer.__init__（vision_transformer.py:60-180）装配的部件：

forward_features_list（vision_transformer.py:222-261）是论文流程的关键：

• RoPE 在每一层都现算（rope_embed(H, W)）→ 把 (H, W) 喂给每个 block。
• 末尾 norm 可以是共享的，也可以分 cls_norm/local_cls_norm（untie_global_and_local_cls_norm 时，训练用第二组 norm 处理 local crop 的 CLS）。

get_intermediate_layers（vision_transformer.py:285-322）：DINOv3 论文强调 dense 特征，这里就是取最后 n 层（或指定层），并支持 reshape=True 把 token 还原为 (B, C, H/p, W/p)。

各种规模（ViT-S/B/L/SO400M/H+/G/7B）在 vision_transformer.py:332-416。

Step 2｜子模块

• PatchEmbed dinov3/layers/patch_embed.py：Conv2d(in_chans, embed_dim, kernel=stride=patch_size)，无 padding，输出 (B, embed_dim, H/P, W/P) 后 flatten。flatten_embedding=False 让主干内部使用 4D 形状。

• 2D RoPE dinov3/layers/rope_position_encoding.py：

  • 坐标归一化到 [-1, 1]（normalize_coords: "separate" 各自归一化），训练时再可选 shift/jitter/rescale 提升分辨率鲁棒性。
  • periods 用 base ** (2i / D_head) 生成（D_head//4 个），与论文公式一致；亦可指定 min_period/max_period。
  • 注意 embed_dim % (4 * num_heads) == 0：每个 head 拿 D_head，分两半分别给 x/y 坐标。

• Self-Attention dinov3/layers/attention.py:43-118：

  • 一个 qkv 线性层，无 bias（qkv_bias=False 是 7B 配置）。
  • compute_attention 拆分 (B, N, 3, H, D/H)，对 q/k 在 rope 之后做 apply_rope，然后 scaled_dot_product_attention（含 flash）。
  • LinearKMaskedBias（mask_k_bias=True 时）：把 K 的 bias 中段 mask 掉，对应论文 7B 模型里的 trick。
  • forward_list 是 cat → 计算 → split，与 block 的 list 模式配套降低 launch overhead。

• SelfAttentionBlock dinov3/layers/block.py:21-212：

  • 标准 Norm → Attn → residual → Norm → FFN → residual。
  • ls1/ls2 = LayerScale（init_values=1e-5），对应论文。
  • sample_drop_ratio = drop_path 在训练时只对 batch 内子集前向、然后用 index_add 加权和回填（Stochastic Depth），列表输入时同样 cat 后处理。

• FFN dinov3/layers/ffn_layers.py：

  • Mlp：标准 Linear → GELU → Drop → Linear。
  • SwiGLUFFN：w1,w2,w3 三投影，中间维度按 2/3 * hidden 再向上对齐 align_to（swiglu32/64/128），用于 7B 节省参数并稳定训练。

• DINOHead dinov3/layers/dino_head.py：3 层 MLP → L2-norm → 线性 (无 bias) 输出 out_dim 个 prototype；DINO 头和 iBOT 头各一个（separate_head=True）。

Step 3｜SSL 训练总装

dinov3/train/ssl_meta_arch.py 的 SSLMetaArch 是 DINOv3 的训练"操作系统"。

构造（ssl_meta_arch.py:35-261）里你会看到三套 backbone+head：

• student（要反向传播）
• teacher（EMA 复制，no grad）
• gram_teacher（可选；要么来自外部 checkpoint、要么从 EMA 周期性地 fork）

forward_backward（ssl_meta_arch.py:355-429）是单个 iteration：

1. 拿到 global_crops (2 张) + local_crops (8 张) + masks（针对 global）；
2. get_teacher_output → teacher 在 global 上做前向，并对 iBOT 选中的掩码 patch 应用 Sinkhorn-Knopp；
3. get_student_output → student 在 global+local 上联合前向；
4. 可选 get_gram_teacher_output（用另一份 teacher 取得 gram 目标）；
5. compute_losses → 求和后 backprop_loss；
6. 主循环在 dinov3/train/train.py:382-580 负责 EMA 更新、gram teacher 更新、checkpoint、eval。

损失组合（ssl_meta_arch.py:584-684）：

• dino_global_crops_loss + dino_local_crops_loss（来自 DINOLoss，按 crop 对数加权）
• koleo_loss（在 student 的 CLS pre-head 上）
• ibot_loss（在 student 掩码 patch 上）
• gram_loss（gram 矩阵 MSE，可选）

EMA（ssl_meta_arch.py:713-726）：teacher = m * teacher + (1-m) * student，用 torch._foreach_mul_/add_ 做高效原地更新。gram_teacher 的更新（ssl_meta_arch.py:728-745）同理但来源是 teacher。

Step 4｜四个损失

Step 5｜ConvNext 分支（蒸馏时用）

dinov3/models/convnext.py 把 4 阶段 ConvNext 包成与 ViT 同样的输出字典（x_norm_clstoken/x_norm_patchtokens），方便 ssl_meta_arch.py 的 _setup_distillation()（ssl_meta_arch.py:263-298）直接复用 heads/loss；最后一层 norm 之前先 mean(-2,-1) 当 CLS。

Step 6｜推理 / hub

dinov3/hub/backbones.py 的 dinov3_vitb16(...) 等工厂：负责下载权重 → build_model_for_eval → FSDP/compile；get_intermediate_layers 是论文"冻结 backbone 取 dense 特征"的标准用法。

3. 论文 ↔ 代码对照表

4. 建议阅读路线（按论文章节）

1. 架构：Step 1+2（先看主干，再下钻到子模块）。
2. 多作物自监督：Step 3，看 forward_backward 与 get_student_output/get_teacher_output。
3. DINO / iBOT / KoLeo 损失：Step 4 前三个文件。
4. Gram Anchoring（DINOv3 新增）：Step 4 的 gram_loss + ssl_meta_arch.py 的 get_gram_teacher_output，配合 dinov3_vit7b16_gram_anchor.yaml 看 gram 段。
5. 训练技巧（EMA、layer-wise LR、freeze last layer、混合精度/FP8）：dinov3/train/train.py、param_groups.py、cosine_lr_scheduler.py。
6. 评估与下游：dinov3/hub/ 下 segmentors/detectors/depthers/classifiers/dinotxt 拼装 backbone 的方式。