크로스 어텐션

# 크로스 어텐션 구현
# pip install torch

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class CrossAttention(nn.Module):
    """
    크로스 어텐션 모듈
    - query_dim: 쿼리 시퀀스의 차원 (디코더 측)
    - context_dim: 컨텍스트 시퀀스의 차원 (인코더 측)
    - heads: 멀티헤드 어텐션의 헤드 수
    """
    def __init__(self, query_dim, context_dim, heads=8, dim_head=64, dropout=0.0):
        super().__init__()
        self.heads = heads
        self.dim_head = dim_head
        inner_dim = heads * dim_head
        self.scale = dim_head ** -0.5

        # Q는 query(디코더)에서, K/V는 context(인코더)에서 생성
        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, query_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x, context):
        """
        x: 쿼리 시퀀스 [batch, seq_len_q, query_dim]
        context: 컨텍스트 시퀀스 [batch, seq_len_ctx, context_dim]
        """
        batch_size = x.shape[0]

        # Q, K, V 계산
        q = self.to_q(x)  # [batch, seq_q, inner_dim]
        k = self.to_k(context)  # [batch, seq_ctx, inner_dim]
        v = self.to_v(context)  # [batch, seq_ctx, inner_dim]

        # 멀티헤드로 reshape
        q = q.view(batch_size, -1, self.heads, self.dim_head).transpose(1, 2)
        k = k.view(batch_size, -1, self.heads, self.dim_head).transpose(1, 2)
        v = v.view(batch_size, -1, self.heads, self.dim_head).transpose(1, 2)

        # 어텐션 스코어 계산
        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        attn = F.softmax(attn, dim=-1)

        # 어텐션 적용
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).reshape(batch_size, -1, self.heads * self.dim_head)

        return self.to_out(out)

# 사용 예시
batch_size = 2
decoder_seq_len = 10  # 디코더 시퀀스 길이
encoder_seq_len = 20  # 인코더 시퀀스 길이
query_dim = 512
context_dim = 768

# 크로스 어텐션 레이어 생성
cross_attn = CrossAttention(query_dim=query_dim, context_dim=context_dim)

# 입력 생성
decoder_hidden = torch.randn(batch_size, decoder_seq_len, query_dim)
encoder_output = torch.randn(batch_size, encoder_seq_len, context_dim)

# 크로스 어텐션 수행
output = cross_attn(decoder_hidden, encoder_output)
print(f"입력 (디코더): {decoder_hidden.shape}")
print(f"컨텍스트 (인코더): {encoder_output.shape}")
print(f"출력: {output.shape}")  # [2, 10, 512]

# 이미지 캡셔닝에서의 크로스 어텐션 활용
# pip install transformers pillow torch

from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image
import requests

# BLIP 모델 로드 (이미지-텍스트 크로스 어텐션 사용)
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

# 샘플 이미지 로드
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

# 이미지 캡션 생성 (내부적으로 크로스 어텐션 사용)
inputs = processor(image, return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=50)

# 결과 디코딩
caption = processor.decode(outputs[0], skip_special_tokens=True)
print(f"생성된 캡션: {caption}")

# 크로스 어텐션 가중치 확인 (선택적)
with torch.no_grad():
    outputs = model(**inputs, output_attentions=True)
    # decoder의 cross_attentions 확인 가능
    if hasattr(outputs, 'cross_attentions') and outputs.cross_attentions:
        print(f"크로스 어텐션 레이어 수: {len(outputs.cross_attentions)}")