DINOv3:一个优秀预训练模型可以做任何视觉任务(复现-持续更新中)

近期Meta开源了DINOv3。

个人认为,相比DINOv2,DINOv3在零样本迁移任务展现出了更强大的能力,比如迁移到了卫星遥感;这对下游任务迁移很重要,毕竟官方数据集中不会包含你的私有数据。迁移到医学相关的领域也会更加方便。

二是,DINOv3能够扩展到更大的图像输入,并且能够避免在特征图中引入额外的嘈杂特征激活。这是我最在意的一点。在预训练模型时,输入的图像往往是较小的尺寸(如224、256、384)这些尺寸。在下游任务适应时,输入图像可能是更大的768、1024、2048,甚至更大的尺寸(比如病理图像、细胞电镜图像)。为了适应这些大尺寸图像输入,通常需要对patch embedding进行插值,使ViT能够适应大尺寸图像,但这种操作就会导致模型精度下降。DINOv3就很好的解决了这一点。

力大砖飞的数据和模型---精巧设计的训练方案

DINOv3把ViT模型参数扩展到了7B大小,把私有数据集扩展到LVD-1689M。在基础模型时代,数据工程已不再是预备步骤,而是研究贡献本身的核心组成部分。

DINOv3通过两个关键步骤Gram Anchoring和引入RoPE来扩展其模型。

尝试一下在自己的数据上预训练

在进行预训练时,需要执行三阶段训练:

       1. DINOv3的预训练

        2.Gram anchoring阶段训练

        3.高分辨率适应训练

训练笔记

正好手上有一部分数据,尝试预训练一个本研究方向的DINOv3模型。

由于官方的数据并不开源,需要自己写一个dataset类。

import os
import numpy as np
from typing import Callable, List, Optional, Tuple, Union
from pathlib import Path
from .extended import ExtendedVisionDataset
from torchvision import datasets, transforms
from PIL import Image
import glob

class MyDINODataset(ExtendedVisionDataset):
    def __init__(
        self,
        root: str,
        extensions=('.tif', '.png', '.TIF', '.jpg', '.tiff'),
        **kwargs
    ):
        super().__init__(root=root, **kwargs)
        
        self.root_path = Path(root)
        
        self.image_paths = []
        for ext in extensions:
            self.image_paths.extend(list(self.root_path.rglob(f'**/*{ext}')))
        
        if not self.image_paths:
            raise FileNotFoundError(f"在目录 {self.root_path} 中未找到任何 {extensions} 格式的图像")

        print(f"无监督数据集初始化完成,共找到 {len(self.image_paths)} 张图像。")

    def __len__(self) -> int:
        return len(self.image_paths)

    def get_image_data(self, index: int) -> bytes:
        image_path = self.image_paths[index]
        with open(image_path, "rb") as f:
            image_data = f.read()
        return image_data

    def get_target(self, index: int) -> Any:
        return None

1.开始第一阶段pretrain训练,为了方便,这里用ViT-Base模型。

2. 训了50000轮,loss不怎么下降了。手动停下,进入第二阶段gram部分。
这里需要注意两个点,一是gram.ckpt需要用eval文件中的ckpt;

CUDA_VISIBLE_DEVICES=1,2,3,4,5,6 PYTHONPATH=${PWD} torchrun --nproc_per_node=6 \
  dinov3/train/train.py \
  --config-file dinov3/configs/train/dinov3_vitbase16_gram_anchor.yaml \
  --output-dir /output_gram/ \
  train.dataset_path=EMCellFound:root="/xx/train_data/" \
  gram.ckpt="/output/eval/training_499999/sharded_teacher_checkpoint/"

二是官方代码中有个小bug,在 ssl_meta_arch.py的485行,应该是
            teacher_patches = backbone_out["x_norm_patchtokens"]

            with torch.no_grad():
                backbone_out = self.gram_teacher.backbone(images, is_training=True)
            # teacher_patches = backbone_out.x_norm_patchtokens  # [n_crops * B, P_T, D]
            teacher_patches = backbone_out["x_norm_patchtokens"]  # [n_crops * B, P_T, D]

这里我pull了一个request,官方也是很快通过了。

也是成功进入了二阶段的训练

-------------------------------------------------

二编    更新日期:2026-1-7

------------------------------------------------

第三阶段训练高分辨率适应训练 high resolution adapt

CUDA_VISIBLE_DEVICES=1,2,3,4,5,6 PYTHONPATH=${PWD} torchrun --nproc_per_node=6 \
  dinov3/train/train.py \
  --config-file dinov3/configs/train/dinov3_vitb_high_res_adapt.yaml \
  --output-dir /xxxxxxx/output/output_higres/ \
  train.dataset_path=EMCellFound:root="/xxxxx/train_data/" \
  gram.ckpt="/xxxxxx/output_gram/ckpt/524999/"  \
  student.resume_from_teacher_chkpt="/xxxxx/output_gram/ckpt/524999/"

训练完后,能够得到一些distcp结尾的模型文件,需要根据torch官方文档将其转化成pth格式。

模型转换

转换代码:

import os

import torch
import torch.distributed.checkpoint as DCP
from torch.distributed.checkpoint.format_utils import dcp_to_torch_save, torch_save_to_dcp

CHECKPOINT_DIR = "/xxxxx/dinov3/output/output_higres/ckpt/29999/"
TORCH_SAVE_CHECKPOINT_DIR = "torch_save_checkpoint.pth"

# convert dcp model to torch.save (assumes checkpoint was generated as above)
dcp_to_torch_save(CHECKPOINT_DIR, TORCH_SAVE_CHECKPOINT_DIR)

但转换完成后,模型中存储了student、teacher、ema_model以及一些训练超参数信息,需要把这里面的权重参数清理一下,只保存teacher模型参数。(这里有一个要注意的点是,不要使用pytorch>=2.6的环境,会无法load预训练模型。可以在其他pytorch<2.6的环境中来运行这个代码和转化模型)

import torch
from collections import OrderedDict

def extract_teacher_backbone(ckpt_path, save_path):
    print(f"Loading checkpoint: {ckpt_path}")
    checkpoint = torch.load(ckpt_path, map_location='cpu')['model']
    
    # checkpoint 是一个保存参数的字典
    state_dict = checkpoint.get('state_dict', checkpoint) 

    new_state_dict = OrderedDict()

    # 我们只关注 'teacher.backbone' 开头的keys
    prefix = 'teacher.backbone.'
    
    print(f"Start extracting keys with prefix: {prefix}")
    
    for k, v in state_dict.items():
        # 1. 只筛选 teacher.backbone
        if k.startswith(prefix):
            # 2. 去掉前缀: teacher.backbone.blocks.0 -> blocks.0
            new_key = k.replace(prefix, '')
            
            # 3. (可选) 如果不需要 RoPE 的缓存 buffer,可以去掉
            if 'rope_embed.periods' in k:
                # 视你的 ViT 实现而定,如果是 standard ViT 通常不需要这个
                # 如果是 RoPE ViT 则需要保留
                pass 
            
            new_state_dict[new_key] = v

    # 保存
    torch.save(new_state_dict, save_path)
    print(f"Extracted {len(new_state_dict)} keys to {save_path}")
    
    # 打印前几个看看格式对不对
    print("Example keys:")
    for i, k in enumerate(new_state_dict.keys()):
        if i < 20: print(k)

# 使用
extract_teacher_backbone("torch_save_checkpoint.pth", "converted_teacher_backbone.pth")

转换好以后得到converted_teacher_backbone.pth文件。然后从dinov3.models.vision_transformer找到对应的模型,比如我这里是vit_base,下面的参数从预训练的config复制过来,一一对应修改。

import torch
from dinov3.models.vision_transformer import vit_base
model = vit_base(
        in_chans= 3,
        ffn_layer= 'swiglu64',
        patch_size= 16,
        drop_path_rate = 0.4,
        layerscale_init= 1.0e-05,
        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,
        pos_embed_rope_base= 100,
        pos_embed_rope_rescale_coords= 2,
        pos_embed_rope_dtype= 'fp32')


state_dict = torch.load("/root/yzy/Code/dinov3/dinov3-1/notebooks/converted_teacher_backbone.pth")
# print(state_dict.keys())
model.load_state_dict(state_dict, strict=True)
model.cuda()

显示下面的情况就算转换成功了,这时你就拥有了自己预训练的dinov3模型。并且可以按照pytorch模型的形式顺利迁移到自己的任务中去。

ClsToken可视化

有了pytorch模型后,我们可以尝试在自己的数据上进行一系列热图可视化。这里我写了一个cls Token可视化的代码,并在我的电镜数据集上进行可视化分析。

读取模型、数据
import pickle
import os
import urllib

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

import torch
import torchvision.transforms.functional as TF
from scipy import signal

import torch
from dinov3.models.vision_transformer import vit_base
model = vit_base(
        in_chans= 3,
        ffn_layer= 'swiglu64',
        patch_size= 16,
        drop_path_rate = 0.4,
        layerscale_init= 1.0e-05,
        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,
        pos_embed_rope_base= 100,
        pos_embed_rope_rescale_coords= 2,
        pos_embed_rope_dtype= 'fp32')


state_dict = torch.load("/root/yzy/Code/dinov3/dinov3-1/notebooks/converted_teacher_backbone.pth")
# print(state_dict.keys())
model.load_state_dict(state_dict, strict=True)
model.cuda()


import numpy as np
image = Image.open("testimg.tif").convert("RGB")
IMAGENET_MEAN = (0.5, 0.5, 0.5)
IMAGENET_STD = (0.5, 0.5, 0.5)
# Resize the image
image = TF.resize(image, (512, 512))
ori_img = image
image = TF.to_tensor(image)
image = TF.normalize(image, mean=IMAGENET_MEAN, std=IMAGENET_STD)

image = image.unsqueeze(0).cuda()
print(image.shape)


模型推理并可视化Cls token
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np

# 假设 image 已经是 (1, 3, H, W) 且 H, W 是 patch_size 的整数倍
# 假设 Patch Size = 14 (DINOv2) 或 16 (DINOv1)
PATCH_SIZE = 16 
H, W = image.shape[2], image.shape[3]
h_feat, w_feat = H // PATCH_SIZE, W // PATCH_SIZE

# 1. 获取输出
# output[-1] 是一个元组: (patch_tokens, cls_token)
# patch_tokens: [B, N, C] -> [1, 4096, 768] (假设输入1024x1024, patch16)
# cls_token:    [B, C]    -> [1, 768]
output = model.get_intermediate_layers(image, n=12, return_class_token=True)
patch_tokens, cls_token = output[-1] 

# 2. 归一化 (为了计算余弦相似度)
patch_tokens = F.normalize(patch_tokens, dim=-1) # [1, N, 768]
cls_token = F.normalize(cls_token, dim=-1)       # [1, 768]

# 3. 计算相似度 (点积)
# cls_token 扩充维度以便广播: [1, 768] -> [1, 768, 1]
# bmm (Batch Matrix Multiplication) 或者简单的 matmul
# 结果 shape: [1, N]
similarity = torch.matmul(patch_tokens, cls_token.unsqueeze(-1)).squeeze()

# 4. 重塑回图像尺寸 (H_feat, W_feat)
similarity_map = similarity.reshape(h_feat, w_feat).detach().cpu().numpy()
print(similarity_map.shape)
# resize to original image size
similarity_map = np.array(Image.fromarray(similarity_map).resize((W, H), resample=Image.BILINEAR))

# paste similarity_map to original image and show

similarity_map_full = Image.blend(Image.fromarray((similarity_map).astype(np.uint8)), ori_img.convert("L"), alpha=0.5)

# 5. 可视化

plt.figure(figsize=(7, 3), dpi=300)
plt.subplot(1, 3, 1)
plt.imshow(similarity_map, cmap='jet') # 或者 'viridis'
plt.axis('off')
plt.title(f"Image")

plt.subplot(1, 3, 2)
plt.imshow(similarity_map_full, cmap='jet')
plt.axis('off')
plt.title(f"Feature")

plt.subplot(1, 3, 3)
plt.imshow(ori_img, cmap='jet')
plt.title(f"Original")
plt.colorbar()
plt.axis('off')
plt.show()

可视化结果:能看到热图高亮部分集中在纹理上。

-----------------------------------------------------------------

后面再更新dinov3在分类、分割等下游任务的代码和应用例子。

更新日期:2026-1-7

评论 7
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

yumaomi

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值