"""SimCLR 的投影头。

对比损失只作用在投影后的空间；聚类时**丢掉投影头、用骨架特征**——
这是 SimCLR 的既有结论：投影头会丢掉对下游任务有用但对比任务用不上的信息。
"""

from __future__ import annotations

import torch
from torch import nn

from models.registry import HEADS


@HEADS.register("projection")
class ProjectionHead(nn.Module):
    def __init__(self, in_channels: int, hidden_channels: int = 256, out_channels: int = 128) -> None:
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.mlp = nn.Sequential(
            nn.Linear(in_channels, hidden_channels),
            nn.BatchNorm1d(hidden_channels),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_channels, out_channels),
        )

    def forward(self, feat: torch.Tensor) -> torch.Tensor:
        if feat.dim() == 4:
            feat = self.pool(feat).flatten(1)
        return self.mlp(feat)