"""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)