modeling_myresnet.py

"""MyResNet 모델 구현.

ResNet (Deep Residual Learning for Image Recognition, He et al., 2015)
논문을 PyTorch + 허깅페이스 transformers 포맷으로 구현한 파일입니다.
"""
from typing import Optional, Union, Tuple

import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import ImageClassifierOutput

from configuration_myresnet import MyResNetConfig


# ============================================================
# Basic Block (ResNet-18/34용) - 논문 Fig 2
# ============================================================
class BasicBlock(nn.Module):
    """2개의 3x3 conv로 이루어진 기본 residual block.

    y = ReLU( BN(conv(ReLU(BN(conv(x))))) + shortcut(x) )
    """

    expansion = 1  # 출력 채널 배수

    def __init__(self, in_channels: int, out_channels: int, stride: int = 1):
        super().__init__()
        # 첫 번째 3x3 conv (stride로 다운샘플링 가능)
        self.conv1 = nn.Conv2d(
            in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False
        )
        self.bn1 = nn.BatchNorm2d(out_channels)

        # 두 번째 3x3 conv (stride=1 고정)
        self.conv2 = nn.Conv2d(
            out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False
        )
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)

        # Shortcut: 차원이 바뀔 때만 projection (1x1 conv) 사용 (논문 Eqn.2, Option B)
        if stride != 1 or in_channels != out_channels * self.expansion:
            self.shortcut = nn.Sequential(
                nn.Conv2d(
                    in_channels,
                    out_channels * self.expansion,
                    kernel_size=1,
                    stride=stride,
                    bias=False,
                ),
                nn.BatchNorm2d(out_channels * self.expansion),
            )
        else:
            self.shortcut = nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        identity = self.shortcut(x)
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out = out + identity           # 논문 Eqn.1: F(x) + x
        return self.relu(out)          # addition 후 ReLU (논문 Sec 3.2)


# ============================================================
# Bottleneck Block (ResNet-50/101/152용) - 논문 Fig 5 오른쪽
# ============================================================
class BottleneckBlock(nn.Module):
    """1x1 -> 3x3 -> 1x1 구조로 채널을 줄였다가 다시 늘리는 bottleneck block.

    마지막 1x1 conv에서 채널을 expansion(=4)배로 확장합니다.
    """

    expansion = 4

    def __init__(self, in_channels: int, out_channels: int, stride: int = 1):
        super().__init__()
        # 1x1 conv: 채널 축소
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)

        # 3x3 conv: 실제 연산 (bottleneck 중심)
        self.conv2 = nn.Conv2d(
            out_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False
        )
        self.bn2 = nn.BatchNorm2d(out_channels)

        # 1x1 conv: 채널 복원 (4배로 확장)
        self.conv3 = nn.Conv2d(
            out_channels, out_channels * self.expansion, kernel_size=1, bias=False
        )
        self.bn3 = nn.BatchNorm2d(out_channels * self.expansion)
        self.relu = nn.ReLU(inplace=True)

        if stride != 1 or in_channels != out_channels * self.expansion:
            self.shortcut = nn.Sequential(
                nn.Conv2d(
                    in_channels,
                    out_channels * self.expansion,
                    kernel_size=1,
                    stride=stride,
                    bias=False,
                ),
                nn.BatchNorm2d(out_channels * self.expansion),
            )
        else:
            self.shortcut = nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        identity = self.shortcut(x)
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        out = out + identity
        return self.relu(out)


# ============================================================
# PreTrainedModel 베이스 클래스
# ============================================================
class MyResNetPreTrainedModel(PreTrainedModel):
    """from_pretrained / save_pretrained 등을 지원하는 베이스 클래스."""

    config_class = MyResNetConfig
    base_model_prefix = "myresnet"
    main_input_name = "pixel_values"
    supports_gradient_checkpointing = False

    def _init_weights(self, module):
        """논문 Sec 3.4: He initialization 사용."""
        if isinstance(module, nn.Conv2d):
            nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
        elif isinstance(module, nn.BatchNorm2d):
            nn.init.constant_(module.weight, 1)
            nn.init.constant_(module.bias, 0)
        elif isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.01)
            if module.bias is not None:
                nn.init.constant_(module.bias, 0)


# ============================================================
# 이미지 분류용 ResNet
# ============================================================
class MyResNetForImageClassification(MyResNetPreTrainedModel):
    """이미지 분류용 ResNet 모델.

    입력: pixel_values (batch, num_channels, H, W)
    출력: ImageClassifierOutput(loss, logits)

    논문 Table 1의 구조를 따릅니다:
        conv1 (7x7, stride=2) -> maxpool -> stage1~4 -> avgpool -> fc
    """

    def __init__(self, config: MyResNetConfig):
        super().__init__(config)
        self.config = config

        # 블록 종류 선택
        block = BasicBlock if config.block_type == "basic" else BottleneckBlock
        self.in_channels = 64

        # ---- Stem (논문 Table 1 conv1) ----
        # 7x7 conv, 64 channels, stride=2 + BN + ReLU + maxpool
        self.stem = nn.Sequential(
            nn.Conv2d(
                config.num_channels, 64, kernel_size=7, stride=2, padding=3, bias=False
            ),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
        )

        # ---- 4 stages (conv2_x ~ conv5_x) ----
        self.stage1 = self._make_stage(
            block, config.hidden_sizes[0], config.layers[0], stride=1
        )
        self.stage2 = self._make_stage(
            block, config.hidden_sizes[1], config.layers[1], stride=2
        )
        self.stage3 = self._make_stage(
            block, config.hidden_sizes[2], config.layers[2], stride=2
        )
        self.stage4 = self._make_stage(
            block, config.hidden_sizes[3], config.layers[3], stride=2
        )

        # ---- Classification head ----
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=1)
        self.classifier = nn.Linear(
            config.hidden_sizes[3] * block.expansion, config.num_labels
        )

        # 가중치 초기화 적용
        self.post_init()

    def _make_stage(self, block, out_channels: int, num_blocks: int, stride: int):
        """하나의 stage(동일 해상도의 블록들)를 만드는 헬퍼.

        첫 블록만 stride로 다운샘플링하고, 나머지는 stride=1.
        """
        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for s in strides:
            layers.append(block(self.in_channels, out_channels, stride=s))
            self.in_channels = out_channels * block.expansion
        return nn.Sequential(*layers)

    def forward(
        self,
        pixel_values: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, ImageClassifierOutput]:
        """순전파.

        Args:
            pixel_values: (batch, num_channels, H, W) 형태의 이미지 텐서.
            labels: (batch,) 형태의 정답 레이블. 주어지면 loss도 함께 반환.
            return_dict: True면 ImageClassifierOutput, False면 튜플 반환.
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # Shape 흐름: (B, 3, 224, 224)
        x = self.stem(pixel_values)    # -> (B, 64, 56, 56)
        x = self.stage1(x)             # -> (B, 64*e, 56, 56)
        x = self.stage2(x)             # -> (B, 128*e, 28, 28)
        x = self.stage3(x)             # -> (B, 256*e, 14, 14)
        x = self.stage4(x)             # -> (B, 512*e, 7, 7)

        x = self.avgpool(x)            # -> (B, 512*e, 1, 1)
        x = torch.flatten(x, 1)        # -> (B, 512*e)
        logits = self.classifier(x)    # -> (B, num_labels)

        # Loss 계산 (Trainer 호환)
        loss = None
        if labels is not None:
            loss_fn = nn.CrossEntropyLoss()
            loss = loss_fn(logits, labels)

        if not return_dict:
            output = (logits,)
            return ((loss,) + output) if loss is not None else output

        return ImageClassifierOutput(loss=loss, logits=logits)