models/google__efficientnet-b0/ds_cfg.py

from transformers import PretrainedConfig
from typing import Literal, Any

__all__ = [
    "BackboneID",
    "BACKBONE_META",
    "BackboneMLPHeadConfig",
]


# ============================================================
# Backbone whitelist + meta registry
# ============================================================

BackboneID = Literal[
    "google/vit-base-patch16-224",
    "microsoft/swin-tiny-patch4-window7-224",
    "microsoft/resnet-50",
    "google/efficientnet-b0",
    "timm/densenet121.tv_in1k",
    "torchvision/densenet121",
]

# ============================================================
# 2) Backbone metadata registry (Feature dim/rule/unfreeze rule)
# 2) 백본 메타 레지스트리 (feature dim/rule/unfreeze rule 고정)
# ============================================================
# This table is the single source of truth for feature extraction and fine-tuning rules per backbone.
# 이 테이블은 backbone별 feature 추출 및 미세조정 규칙의 단일 기준(source of truth)입니다.
#
# The key type is BackboneID to ensure meta keys never drift from the whitelist.
# 키 타입을 BackboneID로 고정하여 메타 키가 화이트리스트와 어긋나지 않게 합니다.
BACKBONE_META: dict[BackboneID, dict[str, Any]] = {
    # -------------------------
    # Transformers (ViT/Swin)
    # -------------------------
    # These backbones come from transformers and typically output hidden states and/or pooler outputs.
    # 이 백본들은 transformers 계열이며 hidden states와 pooler 출력 등을 제공합니다.

    "google/vit-base-patch16-224": {
        # type indicates which loading/forward/extraction pathway the model code should use.
        # type은 모델 코드가 어떤 로딩/forward/feature 추출 경로를 사용할지 결정합니다.
        "type": "vit",

        # feat_dim is the feature vector dimension consumed by the MLP head.
        # feat_dim은 MLP head가 입력으로 받는 feature 벡터 차원입니다.
        "feat_dim": 768,

        # feat_rule defines how to get a (B, feat_dim) tensor from backbone outputs.
        # feat_rule은 backbone 출력에서 (B, feat_dim) 텐서를 얻는 규칙을 정의합니다.
        "feat_rule": "cls",  # Use last_hidden_state[:, 0, :] as CLS token embedding.
                            # last_hidden_state[:, 0, :]를 CLS 토큰 임베딩으로 사용합니다.

        # unfreeze defines the policy to unfreeze layers during stage2 fine-tuning.
        # unfreeze는 stage2 미세조정에서 어떤 레이어를 풀지 정책을 정의합니다.
        "unfreeze": "last_n",  # Unfreeze the last n encoder blocks.
                               # encoder 블록의 마지막 n개를 unfreeze 합니다.

        # has_bn indicates whether BatchNorm exists and should be handled carefully when freezing.
        # has_bn은 BatchNorm 존재 여부이며 freeze 시 특별 취급이 필요한지 판단에 사용합니다.
        "has_bn": False,
    },

    "microsoft/swin-tiny-patch4-window7-224": {
        # This backbone is a Swin Transformer, which may or may not provide a pooler output depending on implementation.
        # 이 백본은 Swin Transformer이며 구현에 따라 pooler output 제공 여부가 달라질 수 있습니다.
        "type": "swin",
        "feat_dim": 768,

        # Prefer pooler output if available, otherwise fall back to mean pooling.
        # pooler가 있으면 우선 사용하고, 없으면 mean pooling으로 대체합니다.
        "feat_rule": "pool_or_mean",

        # Unfreeze strategy is aligned with transformer-style encoder blocks.
        # unfreeze 전략은 transformer 계열 encoder 블록 기준으로 맞춥니다.
        "unfreeze": "last_n",
        "has_bn": False,
    },

    # -------------------------
    # Transformers (CNNs)
    # -------------------------
    # These backbones are CNNs exposed via transformers, usually producing pooled feature vectors or feature maps.
    # 이 백본들은 transformers로 노출된 CNN이며 pooled feature 또는 feature map을 제공합니다.

    "microsoft/resnet-50": {
        # This entry assumes a transformers-compatible ResNet that can expose pooler or a final feature map.
        # 이 항목은 transformers 호환 ResNet이 pooler 또는 최종 feature map을 제공할 수 있다고 가정합니다.
        "type": "resnet",
        "feat_dim": 2048,

        # Use pooler output if the model provides it, otherwise apply global average pooling (GAP).
        # pooler가 있으면 사용하고, 없으면 global average pooling(GAP)을 적용합니다.
        "feat_rule": "pool_or_gap",

        # CNN unfreeze policy can still be expressed as "last_n" at a block/stage granularity in your model code.
        # CNN도 모델 코드에서 block/stage 단위로 last_n 정책을 적용할 수 있습니다.
        "unfreeze": "last_n",
        "has_bn": True,
    },

    "google/efficientnet-b0": {
        # This entry assumes a transformers-compatible EfficientNet that exposes pooled features or a final feature map.
        # 이 항목은 transformers 호환 EfficientNet이 pooled feature 또는 최종 feature map을 제공한다고 가정합니다.
        "type": "efficientnet",
        "feat_dim": 1280,
        "feat_rule": "pool_or_gap",
        "unfreeze": "last_n",
        "has_bn": True,
    },

    # -------------------------
    # timm (DenseNet via HF Hub)
    # -------------------------
    # This backbone is loaded via timm using the "hf_hub:" prefix in your model loader.
    # 이 백본은 모델 로더에서 timm의 "hf_hub:" 프리픽스를 사용해 로드합니다.
    "timm/densenet121.tv_in1k": {
        "type": "timm_densenet",

        # DenseNet-121 final channel dimension is 1024 for the canonical architecture.
        # DenseNet-121의 표준 아키텍처에서 최종 채널 차원은 1024입니다.
        "feat_dim": 1024,

        # timm forward_features typically returns a feature map that you then GAP to (B, C).
        # timm의 forward_features는 보통 feature map을 반환하고 이후 GAP으로 (B, C)를 만듭니다.
        "feat_rule": "timm_gap",

        # DenseNet uses BatchNorm heavily, so freeze_bn behavior matters for stage1/stage2.
        # DenseNet은 BatchNorm 사용이 많아 stage1/stage2에서 freeze_bn 처리가 중요합니다.
        "unfreeze": "last_n",
        "has_bn": True,
    },

    # -------------------------
    # torchvision (DenseNet direct)
    # -------------------------
    # This backbone is intended for torchvision-style loading and feature extraction, not transformers/timm.
    # 이 백본은 transformers/timm이 아니라 torchvision 스타일 로딩 및 feature 추출을 대상으로 합니다.
    "torchvision/densenet121": {
        "type": "torchvision_densenet",
        "feat_dim": 1024,

        # torchvision DenseNet usually exposes .features and you apply GAP to obtain (B, C).
        # torchvision DenseNet은 보통 .features를 노출하며 GAP으로 (B, C)를 얻습니다.
        "feat_rule": "torchvision_densenet_gap",

        # Unfreeze policy remains last_n, but the interpretation must match torchvision module naming.
        # unfreeze 정책은 last_n을 유지하되, 해석은 torchvision 모듈 네이밍에 맞아야 합니다.
        "unfreeze": "last_n",
        "has_bn": True,
    },
}



class BackboneMLPHeadConfig(PretrainedConfig):
    """
    Configuration for Backbone + MLP Head models.

    Backbone + MLP Head 모델을 위한 설정 클래스입니다.
    """

    # This string is used by Hugging Face AutoConfig to identify the config class.
    # Hugging Face AutoConfig가 이 config를 식별하기 위해 사용하는 고유 ID입니다.
    model_type = "backbone-mlphead-224-fixed"

    def __init__(
        self,
        backbone_name_or_path: BackboneID | None = None,
        mlp_head_bottleneck: int = 256,
        mlp_head_dropout: float = 0.2,
        label2id: dict[str, int] | None = None,
        id2label: dict[int, str] | None = None,
        **kwargs,
    ):
        # ============================================================
        # 0) Guard for argument-less construction
        # 0) 무인자 생성 경로 방어
        # ============================================================
        # Transformers may internally construct this config without arguments
        # (e.g., during AutoConfig resolution or Hub loading).
        # Transformers 내부에서 AutoConfig 또는 Hub 로드 과정 중
        # 인자 없이 config를 생성하는 경로가 실제로 존재합니다.
        #
        # In this case, we must NOT validate or raise errors.
        # 이 경우 검증이나 예외를 발생시키면 안 됩니다.
        #
        # The goal here is to provide a minimal, serialization-safe config.
        # 목표는 최소한의 값만 채워 직렬화/역직렬화가 깨지지 않게 하는 것입니다.
        if backbone_name_or_path is None:
            # num_labels may be implicitly assumed by downstream code,
            # so we explicitly set a safe default.
            # num_labels는 downstream 코드에서 암묵적으로 가정되므로
            # 안전한 기본값을 명시적으로 넣어줍니다.
            if "num_labels" not in kwargs:
                kwargs["num_labels"] = 0

            super().__init__(**kwargs)

            # Backbone is intentionally unset in this path.
            # 이 경로에서는 backbone을 의도적으로 설정하지 않습니다.
            self.backbone_name_or_path = None

            # Store MLP head hyperparameters for completeness.
            # MLP head 관련 하이퍼파라미터는 형태 유지를 위해 저장합니다.
            self.mlp_head_bottleneck = int(mlp_head_bottleneck)
            self.mlp_head_dropout = float(mlp_head_dropout)

            # Empty label mappings ensure stable save/load behavior.
            # 빈 label 매핑은 저장/로드 안정성을 보장합니다.
            self.label2id = {}
            self.id2label = {}
            self.num_labels = int(kwargs.get("num_labels", 0))
            return

        # ============================================================
        # 1) Backbone whitelist validation
        # 1) Backbone 화이트리스트 검증
        # ============================================================
        # Only backbones explicitly registered in BACKBONE_META are allowed.
        # BACKBONE_META에 등록된 backbone만 허용합니다.
        #
        # This prevents accidental usage of unsupported or inconsistent backbones.
        # 지원되지 않거나 불일치한 backbone 사용을 원천 차단합니다.
        if backbone_name_or_path not in BACKBONE_META:
            raise ValueError(
                f"Unsupported backbone_name_or_path={backbone_name_or_path}. "
                f"Allowed: {sorted(BACKBONE_META.keys())}"
            )

        # ============================================================
        # 2) Label mapping normalization
        # 2) 라벨 매핑 정규화
        # ============================================================
        # Both label2id and id2label may be None during pure loading scenarios.
        # 단순 로드(from_pretrained) 시 label2id/id2label이 None일 수 있습니다.
        #
        # We allow this here to keep Hub loading and verification stable.
        # Hub 로드 및 검증 안정성을 위해 여기서는 이를 허용합니다.
        #
        # Fail-fast validation should happen at model or training level instead.
        # 실제 사용 검증은 모델 또는 학습 단계에서 fail-fast로 처리해야 합니다.
        if label2id is None and id2label is None:
            # Respect num_labels if explicitly provided, otherwise default to 0.
            # num_labels가 주어졌다면 존중하고, 아니면 0으로 둡니다.
            num_labels = int(kwargs.get("num_labels", 0))
            label2id_norm: dict[str, int] = {}
            id2label_norm: dict[int, str] = {}
        else:
            # If only one mapping is provided, derive the other.
            # 하나만 주어진 경우 나머지를 자동으로 생성합니다.
            if id2label is None:
                id2label = {v: k for k, v in label2id.items()}
            if label2id is None:
                label2id = {v: k for k, v in id2label.items()}

            # Ensure both mappings are consistent in size.
            # 두 매핑의 크기가 일치하는지 확인합니다.
            if len(label2id) != len(id2label):
                raise ValueError(
                    f"label2id/id2label size mismatch: "
                    f"{len(label2id)} vs {len(id2label)}"
                )

            num_labels = len(id2label)
            label2id_norm = dict(label2id)
            id2label_norm = dict(id2label)

        # ============================================================
        # 3) num_labels consistency enforcement
        # 3) num_labels 일관성 강제
        # ============================================================
        # If num_labels is provided via kwargs, it must match inferred labels.
        # kwargs로 num_labels가 들어온 경우 추론된 값과 반드시 일치해야 합니다.
        if "num_labels" in kwargs:
            if (label2id is not None or id2label is not None) and int(kwargs["num_labels"]) != num_labels:
                raise ValueError(
                    f"kwargs['num_labels']={kwargs['num_labels']} "
                    f"!= inferred num_labels={num_labels}"
                )
        else:
            kwargs["num_labels"] = num_labels

        # ============================================================
        # 4) Parent initialization
        # 4) 부모 클래스 초기화
        # ============================================================
        # Initialize PretrainedConfig with normalized label mappings.
        # 정규화된 라벨 매핑을 사용하여 PretrainedConfig를 초기화합니다.
        super().__init__(
            label2id=label2id_norm,
            id2label=id2label_norm,
            **kwargs,
        )

        # ============================================================
        # 5) Explicit attribute assignment for save/load stability
        # 5) 저장/로드 안정성을 위한 명시적 속성 고정
        # ============================================================
        # Explicitly reassign critical fields to avoid subtle serialization issues.
        # 미묘한 직렬화 문제를 방지하기 위해 핵심 필드를 명시적으로 다시 설정합니다.
        self.backbone_name_or_path = backbone_name_or_path
        self.mlp_head_bottleneck = int(mlp_head_bottleneck)
        self.mlp_head_dropout = float(mlp_head_dropout)

        self.label2id = label2id_norm
        self.id2label = id2label_norm
        self.num_labels = int(kwargs["num_labels"])

    def to_dict(self):
        # Call the parent implementation first.
        # 먼저 부모 구현을 호출합니다.
        output = super().to_dict()

        # Force num_labels to be present and consistent.
        # num_labels가 반드시 존재하고 일관되도록 강제합니다.
        output["num_labels"] = int(
            getattr(self, "num_labels", output.get("num_labels", 0))
        )
        return output


# Register this config so it can be resolved via AutoConfig.
# AutoConfig를 통해 이 config가 해석될 수 있도록 등록합니다.
BackboneMLPHeadConfig.register_for_auto_class("AutoConfig")