DINOv2：自监督视觉Transformer

概述

DINOv2是由Meta AI团队在ICLR 2024提出的自监督视觉Transformer¹。它是DINO（Self-Distillation with No Labels）的升级版本，在ImageNet等基准上实现了无需微调即可达到SOTA性能的突破性成果。

版本	年份	骨干网络	ImageNet KNN	关键创新
DINO	2021	ViT-S/16	78.3%	教师-学生蒸馏
DINOv2	2024	ViT-G/14	86.2%	iBOT、patch对齐、多尺度特征

1. 核心思想

1.1 从DINO到DINOv2

DINO的核心是自蒸馏（self-distillation）——使用来自在线网络的知识指导目标网络：

学生网络：接收局部裁剪
教师网络：接收全局裁剪
知识转移：学生预测教师的输出

DINOv2在DINO基础上进行了多项改进：

iBOT预训练策略：在DINO基础上增加masked图像建模
改进的骨干网络：更大的ViT-G、更小的patch大小
多尺度特征：利用不同层的特征用于不同任务
更强的数据管道：更大的预训练数据集

1.2 知识蒸馏框架

DINOv2使用教师-学生架构进行知识蒸馏：

                    教师网络（EMA）
                         ↑
                         │
    ┌────────────────────┴────────────────────┐
    │                                         │
    ↓                                         ↓
全局裁剪 → 学生网络 ← 局部裁剪     全局裁剪 → 教师网络 ← 局部裁剪
    │                                         │
    ↓                                         ↓
  学生输出                                 教师输出
    │                                         │
    └──────────── 软最大化蒸馏 ──────────────┘

2. 数学形式化

2.1 DINO损失

DINO使用软最大化（soft softmax）蒸馏损失：

L_{DINO} = - x \in X_{g} \sum softmax (\frac{g _{θ} ( x )}{τ}) \cdot lo g (\frac{g _{ξ} ( x )}{τ})

其中：

$X_{g}$ 是全局裁剪视图
$g_{θ}$ 是学生网络
$g_{ξ}$ 是教师网络
$τ$ 是温度参数（通常 $τ_{s} = 0.1, τ_{t} = 0.04$ ）

2.2 教师-学生温度

关键设计：学生使用较高温度，教师使用较低温度

p_{s}^{(i)} = \frac{exp ( z _{i} / τ _{s} )}{\sum _{j} exp ( z _{j} / τ _{s} )}, p_{t}^{(i)} = \frac{exp ( z _{i} / τ _{t} )}{\sum _{j} exp ( z _{j} / τ _{t} )}

$p_{s}$ ：学生输出（更平滑）
$p_{t}$ ：教师输出（更尖锐）

2.3 iBOT损失

DINOv2增加了**iBOT（image BERT Online）**损失：

L_{iBOT} = - x \in X_{m} \sum softmax (\frac{z _{x}}{τ}) \cdot lo g (\frac{z ^ _{x}}{τ})

其中：

$X_{m}$ 是被mask的patch tokens
$z_{x}$ 是学生网络对被mask patch的预测
$\overset{z}{^}_{x}$ 是教师网络对原始patch的预测（不mask）

2.4 总损失

L = L_{DINO} + λ L_{iBOT}

其中 $λ$ 通常设置为1。

3. PyTorch实现

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.attention import SDP_FLASH, SDP_EFFICIENT_ATTENTION
import math
 
 
class DINOv2Loss(nn.Module):
    """DINOv2损失函数"""
    def __init__(self, teacher_temp=0.04, student_temp=0.1, 
                 center_momentum=0.999):
        super().__init__()
        self.teacher_temp = teacher_temp
        self.student_temp = student_temp
        self.center_momentum = center_momentum
        self.register_buffer('center', torch.zeros(1, 1, 768))  # ViT-G hidden size
    
    def forward(self, student_output, teacher_output, masks=None):
        """
        Args:
            student_output: [B, N+1, D] 学生输出
            teacher_output: [B, N+1, D] 教师输出
            masks: 可选的mask tensor
        """
        # 提取[CLS] token
        student_cls = student_output[:, 0]  # [B, D]
        teacher_cls = teacher_output[:, 0].detach()  # [B, D]
        
        # 应用中心化
        teacher_cls = teacher_cls - self.center
        
        # 计算softmax概率
        student_prob = F.softmax(student_cls / self.student_temp, dim=-1)
        teacher_prob = F.softmax(teacher_cls / self.teacher_temp, dim=-1)
        
        # 交叉熵损失
        loss = -torch.sum(teacher_prob * torch.log(student_prob + 1e-8), dim=-1)
        loss = loss.mean()
        
        # 更新中心
        batch_center = teacher_cls.mean(dim=0)
        self.center = self.center * self.center_momentum + \
                      batch_center * (1 - self.center_momentum)
        
        return loss
 
 
class MultiCropWrapper(nn.Module):
    """多裁剪数据增强"""
    def __init__(self, backbone, projector):
        super().__init__()
        self.backbone = backbone
        self.projector = projector
    
    def forward(self, x):
        """
        Args:
            x: 输入图像 [B, C, H, W]
        Returns:
            output: 投影后的表示 [B, D]
        """
        # 提取[CLS] token的表示
        outputs = self.backbone(x)
        cls_output = outputs[:, 0]  # [CLS] token
        
        # 投影
        return self.projector(cls_output)
 
 
def dinov2_loss_student_teacher(student_output, teacher_output, 
                                 student_temp=0.1, teacher_temp=0.04,
                                 center=None):
    """
    简化的DINO损失计算
    """
    def SoftmaxKnLoss(out, teacher_out, temp):
        """软最大化蒸馏损失"""
        # 归一化
        out = out - out.mean(dim=-1, keepdim=True)
        teacher_out = teacher_out - teacher_out.mean(dim=-1, keepdim=True)
        
        # 软最大化
        student_log = F.log_softmax(out / temp, dim=-1)
        teacher_prob = F.softmax(teacher_out / temp, dim=-1)
        
        return -(teacher_prob * student_log).sum(dim=-1).mean()
    
    # 只使用全局视图（[CLS] token）
    s_cls = student_output[:, 0]
    t_cls = teacher_output[:, 0].detach()
    
    return SoftmaxKnLoss(s_cls, t_cls, student_temp)
 
 
class iBOTLoss(nn.Module):
    """iBOT损失：masked图像建模"""
    def __init__(self, mask_ratio=0.5, patch_dim=768, eps=1e-6):
        super().__init__()
        self.mask_ratio = mask_ratio
        self.patch_dim = patch_dim
        self.eps = eps
    
    def forward(self, student_output, teacher_output, student_patch_tokens, 
                teacher_patch_tokens, masks):
        """
        Args:
            student_output: 学生对masked patch的预测
            teacher_output: 教师对unmasked patch的预测
            student_patch_tokens: 学生对patch tokens的表示
            teacher_patch_tokens: 教师对patch tokens的表示
            masks: mask indices
        """
        # 获取masked patch的教师表示
        # 只对masked位置计算损失
        masked_t = teacher_patch_tokens[masks].detach()
        
        # 学生预测
        masked_s = student_patch_tokens[masks]
        
        # 对齐维度
        if masked_s.size(-1) != masked_t.size(-1):
            masked_s = F.linear(masked_s, teacher_patch_tokens)
        
        # iBOT损失
        loss = ((masked_s - masked_t) ** 2).mean()
        
        return loss

ViT骨干网络

class ViTBackbone(nn.Module):
    """简化的ViT骨干网络"""
    def __init__(self, image_size=224, patch_size=14, 
                 hidden_size=768, num_layers=12, num_heads=12):
        super().__init__()
        self.patch_size = patch_size
        self.num_patches = (image_size // patch_size) ** 2
        
        # Patch embedding
        self.patch_embed = nn.Conv2d(3, hidden_size, 
                                     kernel_size=patch_size, 
                                     stride=patch_size)
        
        # 可学习[CLS] token
        self.cls_token = nn.Parameter(torch.zeros(1, 1, hidden_size))
        
        # 位置编码
        self.pos_embed = nn.Parameter(
            torch.zeros(1, self.num_patches + 1, hidden_size)
        )
        
        # Transformer blocks
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_size,
            nhead=num_heads,
            dim_feedforward=hidden_size * 4,
            dropout=0.1,
            activation='gelu',
            batch_first=True,
            norm_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
        
        # Layer norm
        self.norm = nn.LayerNorm(hidden_size)
        
        self._init_weights()
    
    def _init_weights(self):
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
    
    def forward(self, x):
        """
        Args:
            x: [B, C, H, W]
        Returns:
            output: [B, N+1, D]
        """
        B = x.size(0)
        
        # Patch embedding
        x = self.patch_embed(x)  # [B, D, H/P, W/P]
        x = x.flatten(2).transpose(1, 2)  # [B, N, D]
        
        # 添加[CLS] token
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)  # [B, N+1, D]
        
        # 添加位置编码
        x = x + self.pos_embed
        
        # Transformer
        x = self.transformer(x)
        x = self.norm(x)
        
        return x

4. DINOv2的核心创新

4.1 iBOT预训练策略

iBOT在DINO基础上增加了一个masked图像建模目标：

随机mask：随机mask掉40-60%的patch tokens
学生预测：学生网络预测被mask位置的表示
教师提供目标：教师网络提供原始表示作为目标

这使得DINOv2学习到更丰富的局部特征。

4.2 Patch对齐

DINOv2显式地对齐不同patch的表示：

def patch_alignment_loss(student_patches, teacher_patches):
    """
    鼓励相同patch位置的表示相似
    """
    return F.mse_loss(student_patches, teacher_patches.detach())

4.3 多尺度特征利用

DINOv2利用Transformer不同层的特征：

层位置	特征类型	适用任务
浅层	局部纹理	分割、检测
中层	语义概念	分类、检索
深层	全局信息	图像级任务

def extract_multiscale_features(model, x):
    """提取多尺度特征"""
    features = []
    x = model.patch_embed(x)
    x = x.flatten(2).transpose(1, 2)
    x = torch.cat([model.cls_token.expand(x.size(0), -1, -1), x], dim=1)
    x = x + model.pos_embed
    
    for i, block in enumerate(model.transformer.blocks):
        x = block(x)
        if i in [3, 6, 9, 11]:  # 选定的层
            features.append(x)
    
    return features

5. 实验结果

5.1 ImageNet分类

方法	ViT-B	ViT-L	ViT-G
DINO	78.3%	80.1%	-
MAE	76.5%	80.6%	-
DINOv2	83.0%	85.7%	86.2%
Supervised	79.9%	82.4%	84.1%

5.2 线性探测

DINOv2的线性探测性能：

数据集	DINOv2 ViT-G	CLIP ViT-L
ImageNet	86.2%	85.3%
CIFAR-100	88.2%	87.6%
Flowers102	96.5%	95.8%
iNat18	81.4%	80.1%

5.3 密集预测任务

DINOv2在密集预测任务上表现优异：

任务	基准	DINOv2性能
语义分割	ADE20K	55.3 mIoU
实例分割	COCO	53.8 AP
深度估计	NYU Depth	0.35 RMSE

6. 与CLIP的对比

特性	DINOv2	CLIP
训练目标	自监督蒸馏	对比学习
训练数据	图像	图文对
零样本能力	有限	强大
特征质量	细粒度高	语义丰富
多模态	否	是
计算成本	较低	较高

7. 使用DINOv2

7.1 模型加载

import torch
 
# 加载预训练模型
dinov2 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitg14')
 
# 或使用更小的模型
dinov2_small = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14')

7.2 特征提取

def extract_dinov2_features(model, images):
    """提取DINOv2特征"""
    with torch.no_grad():
        # 提取所有层的特征
        features = model.get_intermediate_layers(images, n=4)
        
        # 或只获取[CLS] token
        cls_features = model(images)
        
    return features, cls_features
 
 
# 使用示例
from PIL import Image
import torchvision.transforms as transforms
 
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                       std=[0.229, 0.224, 0.225])
])
 
image = Image.open('example.jpg')
x = transform(image).unsqueeze(0)
 
# 提取特征
features = dinov2(x)

7.3 下游任务应用

# 图像检索
def image_retrieval(query_features, gallery_features, top_k=10):
    """基于余弦相似度的图像检索"""
    similarity = F.cosine_similarity(
        query_features.unsqueeze(1), 
        gallery_features.unsqueeze(0), 
        dim=-1
    )
    return similarity.topk(top_k, dim=-1)
 
 
# 语义分割
def segment_with_dinov2(model, image, num_classes):
    """使用DINOv2特征进行语义分割"""
    features = model.get_intermediate_layers(image, n=12)[-1]
    
    # 上采样到原始分辨率
    h, w = image.shape[2:]
    features = F.interpolate(
        features[:, 1:].transpose(1, 2).reshape(1, -1, h//14, w//14),
        size=(h, w),
        mode='bilinear'
    )
    
    # 分割头
    segmentation = nn.Conv2d(features.size(1), num_classes, 1)(features)
    
    return segmentation

8. 实践技巧

8.1 数据增强

DINOv2使用的增强策略：

DINO_AUGMENTATION = {
    'global': transforms.Compose([
        transforms.RandomResizedCrop(224, scale=(0.4, 1.0)),
        transforms.RandomHorizontalFlip(),
        transforms.ColorJitter(0.4, 0.4, 0.4, 0.1),
        transforms.RandomGrayscale(p=0.1),
        transforms.GaussianBlur(kernel_size=11, sigma=(0.1, 2.0)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                           std=[0.229, 0.224, 0.225])
    ]),
    'local': transforms.Compose([
        transforms.RandomResizedCrop(96, scale=(0.05, 0.4)),
        transforms.RandomHorizontalFlip(),
        transforms.ColorJitter(0.4, 0.4, 0.4, 0.1),
        transforms.RandomGrayscale(p=0.1),
        transforms.GaussianBlur(kernel_size=5, sigma=(0.1, 2.0)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                           std=[0.229, 0.224, 0.225])
    ])
}

8.2 超参数配置

参数	推荐值
教师温度 $τ_{t}$	0.04
学生温度 $τ_{s}$	0.1
中心动量	0.999
学习率	0.0005
权重衰减	0.05
批量大小	2048

9. 总结

DINOv2的核心贡献：

无需微调：在大规模预训练后直接用于下游任务
多尺度特征：浅层局部特征 + 深层全局特征
iBOT策略：masked图像建模增强表示质量
SOTA性能：在分类、分割、检索等任务上达到顶尖水平

DINOv2证明了纯自监督学习可以学习到与CLIP相当甚至更好的视觉表示，为视觉基础模型开辟了新道路。

参考

Oquab, M., et al. (2024). “DINOv2: Learning Robust Visual Features Without Supervision”. ICLR 2024. arXiv:2304.07193 ↩

Metaphor

探索

DINOv2：自监督视觉Transformer

概述

1. 核心思想

1.1 从DINO到DINOv2

1.2 知识蒸馏框架

2. 数学形式化

2.1 DINO损失

2.2 教师-学生温度

2.3 iBOT损失

2.4 总损失

3. PyTorch实现

ViT骨干网络

4. DINOv2的核心创新

4.1 iBOT预训练策略

4.2 Patch对齐

4.3 多尺度特征利用

5. 实验结果

5.1 ImageNet分类

5.2 线性探测

5.3 密集预测任务

6. 与CLIP的对比

7. 使用DINOv2

7.1 模型加载

7.2 特征提取

7.3 下游任务应用

8. 实践技巧

8.1 数据增强

8.2 超参数配置

9. 总结

参考

关系图谱

目录

反向链接

Metaphor

探索

DINOv2：自监督视觉Transformer

概述

1. 核心思想

1.1 从DINO到DINOv2

1.2 知识蒸馏框架

2. 数学形式化

2.1 DINO损失

2.2 教师-学生温度

2.3 iBOT损失

2.4 总损失

3. PyTorch实现

ViT骨干网络

4. DINOv2的核心创新

4.1 iBOT预训练策略

4.2 Patch对齐

4.3 多尺度特征利用

5. 实验结果

5.1 ImageNet分类

5.2 线性探测

5.3 密集预测任务

6. 与CLIP的对比

7. 使用DINOv2

7.1 模型加载

7.2 特征提取

7.3 下游任务应用

8. 实践技巧

8.1 数据增强

8.2 超参数配置

9. 总结

参考

Footnotes

关系图谱

目录

反向链接