Siglip2 modeling_siglip2.py 完整实现解析 — 逐类逐变量逐行

1
2
3
4
5
6

@dataclass
class Siglip2VisionOutput(ModelOutput):
    image_embeds: torch.FloatTensor | None = None     # (B, output_dim)
    last_hidden_state: torch.FloatTensor | None = None # (B, N, hidden_size)
    hidden_states: tuple[torch.FloatTensor, ...] | None = None
    attentions: tuple[torch.FloatTensor, ...] | None = None

1
2
3
4
5
6

@dataclass
class Siglip2TextOutput(ModelOutput):
    text_embeds: torch.FloatTensor | None = None      # (B, output_dim)
    last_hidden_state: torch.FloatTensor | None = None # (B, L, hidden_size)
    hidden_states: tuple[torch.FloatTensor, ...] | None = None
    attentions: tuple[torch.FloatTensor, ...] | None = None

1
2
3
4
5
6
7
8
9

@dataclass
class Siglip2Output(ModelOutput):
    loss: torch.FloatTensor | None = None              # (1,)
    logits_per_image: torch.FloatTensor | None = None   # (I, T)
    logits_per_text: torch.FloatTensor | None = None    # (T, I)
    text_embeds: torch.FloatTensor | None = None        # (B, dim)
    image_embeds: torch.FloatTensor | None = None       # (B, dim)
    text_model_output: BaseModelOutputWithPooling = None
    vision_model_output: BaseModelOutputWithPooling = None

1
2

def to_tuple(self) -> tuple[Any]:
    return tuple(v.to_tuple() if isinstance(v, ModelOutput) else v for v in self.values())

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

def __init__(self, config: Siglip2VisionConfig):
    self.config = config
    self.embed_dim = config.hidden_size       # 768 (base)
    self.patch_size = config.patch_size       # 16

    self.patch_embedding = nn.Linear(
        in_features=config.num_channels * self.patch_size * self.patch_size,  # 3*16*16 = 768
        out_features=self.embed_dim,                                          # 768
    )
    self.num_patches = config.num_patches     # 256
    self.position_embedding_size = int(self.num_patches**0.5)  # sqrt(256) = 16
    self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim)

1
2
3
4
5
6

@staticmethod
def resize_positional_embeddings(
    positional_embeddings,  # (grid_h, grid_w, dim) = (16, 16, 768)
    spatial_shapes,         # (B, 2) — 每张图的 (height, width)
    max_length,             # 256 — 最大 patch 数
) -> torch.Tensor:         # (B, max_length, dim)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

# 第 1 步：permute → (1, dim, H, W)
positional_embeddings = positional_embeddings.permute(2, 0, 1).unsqueeze(0)

# 第 2 步：如果设备是 CPU，上转型到 float32
# 原因：CPU 上的 bilinear + antialias 不支持 bfloat16/float16
if positional_embeddings.device.type == "cpu":
    positional_embeddings = positional_embeddings.to(torch.float32)

# 第 3 步：逐样本缩放位置编码
for i in range(batch_size):
    height, width = spatial_shapes[i].tolist()
    # 双线性插值：(16, 16) → (height, width)
    resized_embeddings = F.interpolate(
        positional_embeddings,
        size=(height, width),   # 目标网格尺寸，例如 (4, 6) 表示 24 个 patch
        mode="bilinear",
        align_corners=False,    # 与 TensorFlow 兼容的关键参数
        antialias=True,         # 抗锯齿。对高频位置信息尤其重要
    )
    # reshape: (1, dim, H, W) → (dim, H*W) → (H*W, dim)
    resized_embeddings = resized_embeddings.reshape(embed_dim, height * width).transpose(0, 1)

    # 第 4 步：赋值到预分配张量的对应位置
    resulted_positional_embeddings[i, :height*width] = resized_embeddings
    # 第 5 步：padding 位置用第一个有效位置的编码填充（而非零填充）
    # 为什么是 resized_embeddings[0] 而非 0？
    # 零填充会在后续 attention 计算中产生不连续的位置编码，第一个位置编码提供平滑过渡
    resulted_positional_embeddings[i, height*width:] = resized_embeddings[0]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

def forward(self, pixel_values, spatial_shapes):
    # 1. 线性投影：每个 patch 向量 → embed_dim
    target_dtype = self.patch_embedding.weight.dtype
    patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))

    # 2. 取出位置编码表 + reshape 回 2D → 缩放到样本尺寸 → flatten
    positional_embeddings = self.position_embedding.weight.reshape(
        self.position_embedding_size, self.position_embedding_size, -1
    )
    resized_positional_embeddings = self.resize_positional_embeddings(
        positional_embeddings, spatial_shapes, max_length=pixel_values.shape[1]
    )

    # 3. 相加
    embeddings = patch_embeds + resized_positional_embeddings
    return embeddings

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

def __init__(self, config: Siglip2TextConfig):
    self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
    self.position_embedding = nn.Embedding(config.max_position_embeddings, config.hidden_size)

    # 预计算的 position_ids buffer
    self.register_buffer(
        "position_ids",
        torch.arange(config.max_position_embeddings).expand((1, -1)),
        persistent=False,
    )

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

def forward(self, input_ids, position_ids, inputs_embeds):
    seq_length = input_ids.shape[-1]
    max_position_embedding = self.position_embedding.weight.shape[0]  # 64

    # 长度检查
    if seq_length > max_position_embedding:
        raise ValueError(...)

    # 如果没有显式传入 position_ids，使用预计算的 buf
    if position_ids is None:
        position_ids = self.position_ids[:, :seq_length]

    # 如果没有显式传入 inputs_embeds（如 prompt tuning），从 token id 嵌入
    if inputs_embeds is None:
        inputs_embeds = self.token_embedding(input_ids)

    position_embeddings = self.position_embedding(position_ids)
    embeddings = inputs_embeds + position_embeddings
    return embeddings

1

def eager_attention_forward(module, query, key, value, attention_mask, scaling, dropout=0.0):

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

# 第 1 步：QK^T × scale
attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling

# 第 2 步：加性 mask。-inf 位置的 softmax 输出为 0
if attention_mask is not None:
    attn_weights = attn_weights + attention_mask

# 第 3 步：softmax（关键 —— dtype=torch.float32）
# 为什么 softmax 要上转 float32？
# fp16/bf16 的数值范围有限，softmax 的 exp 容易溢出。
# 上转 float32 计算 softmax，再转回原始 dtype，保证数值稳定
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)

# 第 4 步：dropout（仅在 training=True 时生效）
attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)

# 第 5 步：× V
attn_output = torch.matmul(attn_weights, value)

# 第 6 步：transpose back → contiguous
# 输入: (B, heads, N, dim) → 输出需要: (B, N, heads, dim) → contiguous
# .contiguous() 确保后续 reshape 不会失败
attn_output = attn_output.transpose(1, 2).contiguous()

return attn_output, attn_weights

1
2
3
4
5
6

class Siglip2Attention(nn.Module):
    def __init__(self, config):
        self.embed_dim = config.hidden_size          # 768
        self.num_heads = config.num_attention_heads   # 12
        self.head_dim = self.embed_dim // self.num_heads  # 64
        # 除法检查：768 / 12 = 64，必须整除

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

def forward(self, hidden_states, attention_mask=None):
    # Step 1: 记住输入 shape
    input_shape = hidden_states.shape[:-1]  # (B, N)

    # Step 2: Q/K/V 投影 + reshape 为多头
    # hidden_states: (B, N, 768)
    # → 投影: (B, N, 768)
    # → view(*, -1, head_dim): (B, N, 12, 64)
    # → transpose(1, 2): (B, 12, N, 64)
    hidden_shape = (*input_shape, -1, self.head_dim)
    queries = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    keys    = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    values  = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

    # Step 3: 动态选择注意力实现
    # ALL_ATTENTION_FUNCTIONS 是 HuggingFace 的注意力后端注册表
    # 根据 config._attn_implementation 选择:
    #   "sdpa" → torch.nn.functional.scaled_dot_product_attention
    #   "flash_attention_2" → flash_attn_varlen_func
    #   "eager" → eager_attention_forward (上述手动实现)
    attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
        self.config._attn_implementation, eager_attention_forward
    )

    # Step 4: 调用注意力
    attn_output, attn_weights = attention_interface(
        self, queries, keys, values, attention_mask,
        is_causal=self.is_causal,     # False → 双向注意力
        scaling=self.scale,           # 1/sqrt(64) = 0.125
        dropout=0.0 if not self.training else self.dropout,
    )

    # Step 5: reshape 回 → out_proj
    # (B, 12, N, 64) → (B, N, 768) → Linear → (B, N, 768)
    attn_output = attn_output.reshape(*input_shape, -1).contiguous()
    attn_output = self.out_proj(attn_output)

    return attn_output, attn_weights

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

class Siglip2MLP(nn.Module):
    def __init__(self, config):
        self.activation_fn = ACT2FN[config.hidden_act]     # "gelu_pytorch_tanh" → GELU(tanh approx)
        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)  # 768 → 3072
        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)  # 3072 → 768

    def forward(self, hidden_states):
        hidden_states = self.fc1(hidden_states)       # 扩展 4x
        hidden_states = self.activation_fn(hidden_states)  # GELU
        hidden_states = self.fc2(hidden_states)       # 压缩回原尺寸
        return hidden_states

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

class Siglip2EncoderLayer(GradientCheckpointingLayer):
    def __init__(self, config):
        self.layer_norm1 = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
        self.self_attn = Siglip2Attention(config)
        self.layer_norm2 = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
        self.mlp = Siglip2MLP(config)

    def forward(self, hidden_states, attention_mask):
        # Pre-LN 架构：先 Norm，后 Attention/MLP
        # 子层 1: Self-Attention
        residual = hidden_states
        hidden_states = self.layer_norm1(hidden_states)
        hidden_states, _ = self.self_attn(hidden_states, attention_mask)
        hidden_states = residual + hidden_states

        # 子层 2: MLP
        residual = hidden_states
        hidden_states = self.layer_norm2(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

class Siglip2Encoder(nn.Module):
    def __init__(self, config):
        self.layers = nn.ModuleList([
            Siglip2EncoderLayer(config) for _ in range(config.num_hidden_layers)
        ])
        self.gradient_checkpointing = False

    def forward(self, inputs_embeds, attention_mask):
        hidden_states = inputs_embeds
        for encoder_layer in self.layers:
            hidden_states = encoder_layer(hidden_states, attention_mask)
        return BaseModelOutput(last_hidden_state=hidden_states)

1
2
3
4
5

class Siglip2PreTrainedModel(PreTrainedModel):
    config: Siglip2Config
    base_model_prefix = "siglip2"
    input_modalities = ("image", "text")
    supports_gradient_checkpointing = True

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

def _init_weights(self, module):
    super()._init_weights(module)  # 先调用基类的默认初始化

    if isinstance(module, Siglip2VisionEmbeddings):
        # 位置编码: N(0, 1/sqrt(width))
        init.normal_(module.position_embedding.weight, std=1 / np.sqrt(width))

    elif isinstance(module, nn.Embedding):
        # 文本 token/position embedding: Flax 默认截断正态初始化
        init.default_flax_embed_init_(module.weight)

    elif isinstance(module, Siglip2Attention):
        # Q/K/V/O 投影: Xavier Uniform
        init.xavier_uniform_(module.q_proj.weight)
        init.xavier_uniform_(module.k_proj.weight)
        init.xavier_uniform_(module.v_proj.weight)
        init.xavier_uniform_(module.out_proj.weight)
        # bias 初始化为 0
        init.zeros_(module.q_proj.bias)
        init.zeros_(module.k_proj.bias)
        init.zeros_(module.v_proj.bias)
        init.zeros_(module.out_proj.bias)

    elif isinstance(module, Siglip2MLP):
        # FC1/FC2: Xavier Uniform
        init.xavier_uniform_(module.fc1.weight)
        init.xavier_uniform_(module.fc2.weight)
        # bias: N(0, 1e-6) — 极小值，几乎为 0 但有微小扰动防止对称性
        init.normal_(module.fc1.bias, std=1e-6)
        init.normal_(module.fc2.bias, std=1e-6)

    elif isinstance(module, Siglip2MultiheadAttentionPoolingHead):
        # probe (可学习查询): Xavier Uniform
        init.xavier_uniform_(module.probe)
        # MultiheadAttention 内置权重: Xavier Uniform
        init.xavier_uniform_(module.attention.in_proj_weight)
        init.zeros_(module.attention.in_proj_bias)

    elif isinstance(module, Siglip2Model):
        # logit_scale 和 logit_bias 初始化为 0
        # 训练开始时 logits = 0 * exp(0) + 0 = 0 → sigmoid(0) = 0.5
        # 即初始时模型对任何 image-text pair 都输出 0.5 概率
        init.zeros_(module.logit_scale)
        init.zeros_(module.logit_bias)

    elif isinstance(module, Siglip2ForImageClassification):
        # 分类头: N(0, 1/sqrt(hidden_size) * initializer_factor)
        init.normal_(module.classifier.weight,
                     std=config.hidden_size**-0.5 * config.initializer_factor)

    elif isinstance(module, (nn.Linear, nn.Conv2d)):
        # 所有其他线性层/卷积: LeCun 正态初始化
        init.lecun_normal_(module.weight)
        if module.bias is not None:
            init.zeros_(module.bias)

1
2
3
4
5
6
7
8

def __init__(self, config: Siglip2VisionConfig):
    self.embeddings = Siglip2VisionEmbeddings(config)
    self.encoder = Siglip2Encoder(config)
    self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
    self.use_head = True if not hasattr(config, "vision_use_head") else config.vision_use_head
    if self.use_head:
        self.head = Siglip2MultiheadAttentionPoolingHead(config)
    self.post_init()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

def forward(self, pixel_values, pixel_attention_mask, spatial_shapes):
    # 第 1 步：Embedding（线性投影 + 位置编码）
    hidden_states = self.embeddings(pixel_values, spatial_shapes)

    # 第 2 步：创建双向 attention mask
    # create_bidirectional_mask 将 1D mask (B, N) 展开为 (B, 1, N, N) 的加性 mask
    # 有效位置 = 0, padding 位置 = -inf
    encoder_attention_mask = create_bidirectional_mask(
        config=self.config,
        inputs_embeds=hidden_states,
        attention_mask=pixel_attention_mask,
    )

    # 第 3 步：Transformer 编码
    encoder_outputs = self.encoder(
        inputs_embeds=hidden_states,
        attention_mask=encoder_attention_mask,
    )

    # 第 4 步：后 LayerNorm
    last_hidden_state = self.post_layernorm(encoder_outputs.last_hidden_state)

    # 第 5 步：多头注意力池化
    pooler_output = self.head(last_hidden_state, pixel_attention_mask) if self.use_head else None

    return BaseModelOutputWithPooling(
        last_hidden_state=last_hidden_state,
        pooler_output=pooler_output,
    )

1
2
3
4
5
6
7
8
9

# 输入: pixel_attention_mask = (B, N)，值为 0 或 1
# 输出: (B, 1, N, N)，格式:
#   mask[i, j] = 0     if pixel_attention_mask[i] == 1 and pixel_attention_mask[j] == 1
#   mask[i, j] = -inf  otherwise
#
# 加性 mask 在 softmax 之前与 QK^T 相加:
#   attn = softmax(QK^T/√d + mask)
# 有效位置: 加上 0，不影响
# padding 位置: 加上 -inf，softmax 后 → 0

1
2
3
4
5
6

def __init__(self, config: Siglip2TextConfig):
    self.embeddings = Siglip2TextEmbeddings(config)
    self.encoder = Siglip2Encoder(config)
    self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
    self.head = nn.Linear(embed_dim, config.projection_size)
    self.post_init()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

def forward(self, input_ids, attention_mask, position_ids):
    # 第 1 步：文本 Embedding
    hidden_states = self.embeddings(input_ids, position_ids)

    # 第 2 步：创建双向 attention mask（非因果！）
    attention_mask = create_bidirectional_mask(
        config=self.config,
        inputs_embeds=hidden_states,
        attention_mask=attention_mask,
    )

    # 第 3 步：Transformer 编码
    encoder_outputs = self.encoder(hidden_states, attention_mask)

    # 第 4 步：后 LayerNorm
    last_hidden_state = self.final_layer_norm(encoder_outputs.last_hidden_state)

    # 第 5 步：EOS token 池化 + 投影
    pooled_output = last_hidden_state[:, -1, :]      # 取最后一个 token
    pooled_output = self.head(pooled_output)          # 线性投影

    return BaseModelOutputWithPooling(
        last_hidden_state=last_hidden_state,
        pooler_output=pooled_output,
    )

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

class Siglip2MultiheadAttentionPoolingHead(nn.Module):
    def __init__(self, config):
        # 可学习的 probe 向量（类比 CLS token）
        self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))

        # 多头注意力：probe 作为 query，patch tokens 作为 key/value
        self.attention = nn.MultiheadAttention(
            config.hidden_size, config.num_attention_heads, batch_first=True
        )

        # 后处理：LayerNorm + MLP (残差连接)
        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.mlp = Siglip2MLP(config)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

if attention_mask is not None:
    # 步骤 3a: 创建 (B*heads, target_len, source_len) 的 mask
    attention_mask = attention_mask.repeat(1, self.num_heads, target_len, 1)
    attention_mask = attention_mask.reshape(-1, target_len, source_len)

    # 步骤 3b: 布尔 mask → 加性 mask
    if attention_mask.dtype == torch.bool:
        attention_mask = torch.where(
            attention_mask,
            torch.tensor(0.0, device=attention_mask.device, dtype=probe.dtype),
            torch.finfo(probe.dtype).min,  # ≈ -65504 for fp16
        )

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

def __init__(self, config: Siglip2Config):
    text_config = config.text_config
    vision_config = config.vision_config

    # 通过 _from_config 创建子模型（确保注意力配置传递）
    self.text_model = Siglip2TextModel._from_config(text_config)
    self.vision_model = Siglip2VisionModel._from_config(vision_config)

    # 可学习的 logit_scale 和 logit_bias
    self.logit_scale = nn.Parameter(torch.randn(1))
    self.logit_bias = nn.Parameter(torch.randn(1))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

def forward(self, input_ids, pixel_values, pixel_attention_mask, spatial_shapes,
            attention_mask, position_ids, return_loss):
    # ===== 第 1 步：编码图像和文本 =====
    vision_outputs = self.vision_model(
        pixel_values=pixel_values,
        pixel_attention_mask=pixel_attention_mask,
        spatial_shapes=spatial_shapes,
    )
    text_outputs = self.text_model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        position_ids=position_ids,
    )

    image_embeds = vision_outputs.pooler_output   # (I, dim)
    text_embeds = text_outputs.pooler_output      # (T, dim)

    # ===== 第 2 步：L2 归一化 =====
    image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
    text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)

1
2
3

    # ===== 第 3 步：计算余弦相似度矩阵 =====
    logits_per_text = torch.matmul(text_embeds, image_embeds.t().to(text_embeds.device))
    # shape: (T, I) = (文本数, 图像数)

1
2
3
4

    # ===== 第 4 步：应用可学习温度和偏置 =====
    logit_scale = self.logit_scale.to(text_embeds.device)
    logit_bias = self.logit_bias.to(text_embeds.device)
    logits_per_text = logits_per_text * logit_scale.exp() + logit_bias

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

    # ===== 第 5 步：计算 Sigmoid Loss =====
    loss = None
    if return_loss:
        # 构造标签矩阵: 对角线=1 (正样本), 非对角线=-1 (负样本)
        eye = torch.eye(logits_per_text.size(0), device=logits_per_text.device)  # (T, I) 对角线
        m1_diag1 = -torch.ones_like(logits_per_text) + 2 * eye
        # 对角线: -1 + 2*1 = +1
        # 非对角线: -1 + 2*0 = -1

        # log-sigmoid: log(1/(1+exp(-x))) = -softplus(-x)
        loglik = torch.nn.functional.logsigmoid(m1_diag1 * logits_per_text)
        # 正样本: logsigmoid(+1 * logit) → loglik 高 = loss 低
        # 负样本: logsigmoid(-1 * logit) → loglik 低 = loss 高

        nll = -torch.sum(loglik, dim=-1)  # 对图像维度求和
        loss = nll.mean()                  # 对 batch 求平均

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

eye = torch.eye(T, device=...)       # [[1,0,0], [0,1,0], [0,0,1]]
m1_diag1 = -torch.ones_like(T) + 2*eye

# 结果:
# [[+1, -1, -1],
#  [-1, +1, -1],
#  [-1, -1, +1]]
#
# 第 (i,j) 个元素:
#   i==j → +1 (正样本 target)
#   i!=j → -1 (负样本 target)

1
2
3
4
5

def get_text_features(self, input_ids, attention_mask, position_ids):
    return self.text_model(input_ids, attention_mask, position_ids)

def get_image_features(self, pixel_values, pixel_attention_mask, spatial_shapes):
    return self.vision_model(pixel_values, pixel_attention_mask, spatial_shapes)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

class Siglip2ForImageClassification(Siglip2PreTrainedModel):
    def __init__(self, config):
        self.vision_model = Siglip2VisionModel._from_config(config.vision_config)
        self.classifier = (
            nn.Linear(config.vision_config.hidden_size, config.num_labels)
            if config.num_labels > 0 else nn.Identity()
        )

    def forward(self, pixel_values, pixel_attention_mask, spatial_shapes, labels):
        outputs = self.vision_model(pixel_values, pixel_attention_mask, spatial_shapes)
        sequence_output = outputs.last_hidden_state  # (B, N, 768)

        # 平均池化（而非 MultiheadAttentionPoolingHead）
        if pixel_attention_mask is not None:
            pool_mask = pixel_attention_mask[..., None].to(sequence_output.device)  # (B, N, 1)
            sequence_output = torch.sum(sequence_output * pool_mask, dim=1) / torch.sum(pool_mask, dim=1)
        else:
            sequence_output = torch.mean(sequence_output, dim=1)

        logits = self.classifier(sequence_output)
        loss = self.loss_function(labels, logits, self.config) if labels is not None else None

        return ImageClassifierOutput(loss=loss, logits=logits)