Boostcamp AI Tech (Day 022)

1. 과제3

CLIP: supervised dataset을 이용하여 image-text pair에 대해 contrastive learning으로 학습하는 방법. image encoder와 text encoder가 있다. 각 encoder에서 나온 feature들이 positive pair일 경우는 similarity가 커지도록, negative pair일 경우는 similarity가 작아지도록 학습된다.

pretrained CLIP을 이용하여 모델이 이미지의 어떤 부분을 참조하여 feature를 생성하고 있는지를 visualize하여 확인해보자! Attention을 visualize하므로, image encoder는 ViT를 사용한다.

Model Architectures

1. Bottle Neck 클래스 ResNet의 기본 블록인 Bottleneck은 Residual Learning을 위해 사용된다. Residual 블록은 입력을 여러 레이어를 통해 변환한 출력에 입력을 더하는 구조를 가지고 있다.

2. AttentionPool2d 클래스 이미지를 인코딩한 후 Attention Pooling 레이어를 통해 각 패치의 중요도를 계산한다. 이미지 전체에 대한 맥락적 정보를 추출한다.

3. ModifiedResNet 클래스 이 클래스는 ResNet의 변형 버전으로, CLIP 모델의 이미지 인코더로 사용된다. 3개의 “stem” convolution을 사용하며, 각 레이어에서는 anti-aliasing을 위해 avgpool을 사용한다.

class CustomBottleneck(nn.Module):
    expansion_factor = 4  # 확장 비율

    def __init__(self, input_channels, output_channels, stride=1):
        super(CustomBottleneck, self).__init__()

1.Bottle Neck

1.1 입력과 출력의 연결 (Residual Connection)

• 네트워크의 핵심 아이디어는 입력을 여러 레이어를 거쳐 변형한 후, 변형된 출력에 입력을 더하는 것이다.
• 네트워크가 깊어지더라도 기울기 소실(vanishing gradient) 문제를 완화한다.

1.2 컨볼루션 연산

• Conv1: 첫 번째 1x1 컨볼루션은 입력 채널의 수를 줄이거나 늘린다. 정보의 압축 및 확장.
• Conv2: 두 번째 3x3 컨볼루션은 이미지의 공간 정보를 유지하면서 특징을 추출한다.
• Conv3: 마지막 1x1 컨볼루션은 채널 수를 다시 확장하며, expansion_factor에 의해 확장된 채널 수를 만든다.

        # 첫 번째 1x1 컨볼루션
        self.conv1 = nn.Conv2d(input_channels, output_channels, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(output_channels)
        self.relu = nn.ReLU(inplace=True)  # 모든 ReLU에 동일한 인스턴스를 사용

        # 두 번째 3x3 컨볼루션
        self.conv2 = nn.Conv2d(output_channels, output_channels, kernel_size=3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(output_channels)

        # 필요 시 다운샘플링을 위한 평균 풀링 레이어
        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()

        # 마지막 1x1 컨볼루션
        self.conv3 = nn.Conv2d(output_channels, output_channels * self.expansion_factor, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(output_channels * self.expansion_factor)

        # 다운샘플링 레이어, 입력과 출력 크기가 다를 경우에만 사용
        self.downsample = None
        if stride > 1 or input_channels != output_channels * self.expansion_factor:
            self.downsample = nn.Sequential(OrderedDict([
                ("avgpool", nn.AvgPool2d(stride)),
                ("conv", nn.Conv2d(input_channels, output_channels * self.expansion_factor, kernel_size=1, bias=False)),
                ("bn", nn.BatchNorm2d(output_channels * self.expansion_factor))
            ]))

1.3 다운샘플링

• 중간 레이어에서 스트라이드가 2 이상인 경우, 평균 풀링(AvgPool2d)을 사용하여 공간 크기를 절반으로 줄인다.
• 입력과 출력의 크기가 다를 경우엔 별도의 다운샘플링 레이어가 추가되어 입력과 출력의 크기를 맞춘다.

      def forward(self, x: torch.Tensor):
     # Identity 연결을 위한 입력 보존
     identity = x
  
     # 첫 번째 1x1 컨볼루션 및 활성화
     out = self.relu(self.bn1(self.conv1(x)))
     # 두 번째 3x3 컨볼루션 및 활성화
     out = self.relu(self.bn2(self.conv2(out)))
     # 필요 시 다운샘플링
     out = self.avgpool(out)
     # 마지막 1x1 컨볼루션 및 배치 정규화
     out = self.bn3(self.conv3(out))
  
     # 입력과 크기를 맞추기 위한 다운샘플링 처리
     if self.downsample is not None:
         identity = self.downsample(x)
  
     # 출력과 입력을 더해 Residual Connection
     out += identity
     # 최종 활성화
     out = self.relu(out)
     return out
  

1.4 배치 정규화 및 활성화 함수

• 각 컨볼루션 뒤에는 배치 정규화(BatchNorm2d)가 위치하며, 이를 통해 각 배치의 통계적 변동을 줄여 안정적인 학습이 가능하다.
• ReLU 활성화 함수는 비선형성을 부여하여 모델의 표현력을 증가시킨다.

1.5 Residual Connection의 장점

• Residual Connection은 입력 데이터가 변형된 출력 데이터와 더해짐으로써, 모델이 학습 초기 단계에서도 유용한 특징을 유지할 수 있도록 도와준다.
• 깊은 네트워크에서의 중요한 특징이다.

2. ResidualAttentionBlock

트랜스포머(Transformer) 모델의 주요 구성 요소인 어텐션(attention) 메커니즘과 잔차 연결(residual connection)을 결합한 블록. 이 블록은 입력을 받아 어텐션을 수행하고, 그 결과를 원래 입력에 더해준 후, 다시 MLP를 거쳐 최종 출력을 생성한다.

class ResidualAttentionBlock(nn.Module):
    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
        super().__init__()

        self.attn = nn.MultiheadAttention(d_model, n_head)
        self.ln_1 = LayerNorm(d_model)  # 첫 번째 LayerNorm
        self.mlp = nn.Sequential(OrderedDict([
            ("c_fc", nn.Linear(d_model, d_model * 4)),  # 입력 차원을 4배로 확장
            ("gelu", QuickGELU()),  # GELU 활성화 함수
            ("c_proj", nn.Linear(d_model * 4, d_model))  # 다시 원래 차원으로 축소
        ]))
        self.ln_2 = LayerNorm(d_model)  # 두 번째 LayerNorm
        self.attn_mask = attn_mask

2.1 어텐션 함수 (attention)

    def attention(self, x: torch.Tensor):
        # 어텐션 마스크가 있으면, 해당 마스크를 사용하여 어텐션 계산
        if self.attn_mask is not None:
            self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device)
        # 멀티헤드 어텐션을 수행하고, 어텐션 결과와 가중치를 반환
        attn_output, attn_weights = self.attn(x, x, x, need_weights=True, attn_mask=self.attn_mask, average_attn_weights=False)
        return attn_output, attn_weights

• input tensor x에 대해 self.attn을 통해 Multi-head attention을 수행한다. 특정 위치의 attention을 제한할 수 있다.
• 어텐션의 출력 (attn_output)과 어텐션 가중치 (attn_weights)를 반환한다.

2.2 forward 함수

    def forward(self, x: torch.Tensor):
        # 잔차 연결 및 첫 번째 LayerNorm을 통한 어텐션 계산
        attention_res, weights = self.attention(self.ln_1(x))
        x = x + attention_res  # 잔차 연결

        # 잔차 연결 및 두 번째 LayerNorm을 통한 MLP 계산
        x = x + self.mlp(self.ln_2(x))  # 잔차 연결
        return x, weights

• 첫 번째 LayerNorm을 거친 후, 어텐션을 계산한다. 계산된 어텐션 결과는 잔차 연결(residual connection)을 통해 원래의 입력에 concat된다.
• 두 번째 LayerNorm을 거쳐 MLP를 통과시키고, 그 결과를 다시 잔차 연결을 통해 원래의 입력 concat해 최종 출력.
• 어텐션의 결과& 가중치를 반환