输入图像         [B, 3, 224, 224]
     │
     ├─ PatchEmbed (Conv2d 16×16, stride 16)
     ▼
patch tokens     [B, 14, 14, 384]   ← H=224/16=14, W=14
     │ flatten(1,2)
     ▼
patch tokens     [B, 196, 384]      ← 196 个 patch
     │
     ├─ cat(cls_token)              # 1 个
     ├─ cat(storage_tokens × 4)     # 4 个（DINOv3 特有）
     ▼
完整序列         [B, 1+4+196, 384] = [B, 201, 384]
     │
     ├─ 12 × SelfAttentionBlock (RoPE + LayerScale)
     ▼
最终 token 序列   [B, 201, 384]
     │
     ├─ 切分：[:, 0] → CLS, [:, 1:5] → storage, [:, 5:] → patches
     └─ self.norm(...)              # 最后再过一层 LayerNorm

from dinov3.hub.backbones import dinov3_vits16
import torch

backbone = dinov3_vits16(pretrained=True).eval()
x = torch.randn(1, 3, 224, 224)

# 方式 A：只取 patch tokens 的空间特征
feat = backbone.get_intermediate_layers(x, n=1, reshape=True)[0]
# feat.shape == [1, 384, 14, 14]   ← BCHW，stride=16

(patches, cls, storage) = backbone.get_intermediate_layers(
    x, n=1, reshape=True, return_class_token=True, return_extra_tokens=True
)[0]
# patches:  [1, 384, 14, 14]
# cls:      [1, 384]
# storage:  [1, 4, 384]

# 取 4 个 block 的输出
feats = backbone.get_intermediate_layers(
    x, n=[3, 6, 9, 11],   # 0-indexed：第 4、7、10、12 个 block
    reshape=True,
    norm=True,            # 每个特征都过最终的 self.norm
)
for i, f in enumerate(feats):
    print(i, f.shape)
# 0 torch.Size([1, 384, 14, 14])   ← stride 16, 浅层
# 1 torch.Size([1, 384, 14, 14])   ← stride 16
# 2 torch.Size([1, 384, 14, 14])   ← stride 16
# 3 torch.Size([1, 384, 14, 14])   ← stride 16, 最深层

import torch.nn as nn

class ViTMultiScale(nn.Module):
    def __init__(self, backbone, embed_dim=384):
        super().__init__()
        self.backbone = backbone
        # 4 个 1/16 特征 → 4 个不同上/下采样倍率
        self.up_4x = nn.Sequential(           # 1/16 → 1/4
            nn.ConvTranspose2d(embed_dim, embed_dim//2, 4, 2, 1),
            nn.ConvTranspose2d(embed_dim//2, embed_dim//4, 4, 2, 1),
            nn.GELU(),
            nn.Conv2d(embed_dim//4, 256, 3, 1, 1),
        )
        self.up_2x = nn.Sequential(           # 1/16 → 1/8
            nn.ConvTranspose2d(embed_dim, embed_dim//2, 4, 2, 1),
            nn.GELU(),
            nn.Conv2d(embed_dim//2, 256, 3, 1, 1),
        )
        self.up_1x = nn.Sequential(           # 1/16（保持）
            nn.Conv2d(embed_dim, 256, 3, 1, 1),
        )
        self.down_2x = nn.Sequential(         # 1/16 → 1/32
            nn.Conv2d(embed_dim, 256, 3, 2, 1),  # stride 2
            nn.GELU(),
            nn.Conv2d(256, 256, 3, 1, 1),
        )

    def forward(self, x):
        feats = self.backbone.get_intermediate_layers(
            x, n=[3, 6, 9, 11], reshape=True, norm=True
        )
        f_shallow, f_mid1, f_mid2, f_deep = feats

        return {
            "stride_4":  self.up_4x(f_shallow),    # [B, 256, 56, 56]
            "stride_8":  self.up_2x(f_mid1),       # [B, 256, 28, 28]
            "stride_16": self.up_1x(f_mid2),       # [B, 256, 14, 14]
            "stride_32": self.down_2x(f_deep),     # [B, 256, 7, 7]
        }

x = torch.randn(1, 3, 512, 768)   # 任意尺寸都行
feat = backbone.get_intermediate_layers(x, n=1, reshape=True)[0]
# feat.shape == [1, 384, 32, 48]   ← H/16, W/16

backbone = backbone.to(memory_format=torch.channels_last)
backbone = backbone.half()  # 或 .bfloat16()