Transformer가 시대를 바꾼다.

import torch
import torch.nn as nn
import math

# 1. Scaled Dot-Product Attention
class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super().__init__()
        
    def forward(self, query, key, value, mask=None):
        # query, key, value: Attention 메커니즘의 입력 텐서들
        # mask: 특정 위치를 무시하기 위해 사용하는 텐서 (예: 패딩 위치) 😊

        # 어텐션 스코어 계산: Q x K^T / sqrt(d_k)
        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(query.size(-1))
        
        if mask is not None:
            # 패딩 마스크 위치를 매우 작은 값(-∞)으로 설정해 softmax에서 제외
            scores = scores.masked_fill(mask == 0, -1e9)
        
        # 소프트맥스를 통해 확률 분포 계산 (어텐션 가중치)
        attention_weights = torch.softmax(scores, dim=-1)

        # 어텐션 결과 = 가중치 x V
        return torch.matmul(attention_weights, value), attention_weights

# 2. Multi-Head Attention
class MultiHeadAttention(nn.Module):
    def __init__(self, embed_size, num_heads):
        super().__init__()
        self.num_heads = num_heads  # 헤드 개수
        self.head_dim = embed_size // num_heads  # 각 헤드의 차원
        
        # 임베딩 크기가 헤드 수로 나누어 떨어지는지 확인
        assert self.head_dim * num_heads == embed_size, "Embedding size must be divisible by number of heads"
        
        # 쿼리(Q), 키(K), 값(V)를 생성하는 선형 변환
        self.query = nn.Linear(embed_size, embed_size)
        self.key = nn.Linear(embed_size, embed_size)
        self.value = nn.Linear(embed_size, embed_size)
        
        # 최종 출력 선형 변환
        self.fc_out = nn.Linear(embed_size, embed_size)
        
    def forward(self, query, key, value, mask=None):
        N = query.shape[0]  # 배치 크기
        
        # Q, K, V를 각각 (배치, 헤드 수, 문장 길이, 헤드 차원)으로 변환
        Q = self.query(query).view(N, -1, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.key(key).view(N, -1, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.value(value).view(N, -1, self.num_heads, self.head_dim).transpose(1, 2)
        
        # Scaled Dot-Product Attention 수행
        attention, _ = ScaledDotProductAttention()(Q, K, V, mask)
        
        # 어텐션 결과를 결합하고 선형 변환
        attention = attention.transpose(1, 2).contiguous().view(N, -1, self.num_heads * self.head_dim)
        return self.fc_out(attention)

# 3. Feedforward Neural Network
class FeedForward(nn.Module):
    def __init__(self, embed_size, ff_hidden_size):
        super().__init__()
        # 첫 번째 선형 계층: 임베딩 크기를 확장
        self.fc1 = nn.Linear(embed_size, ff_hidden_size)
        # 두 번째 선형 계층: 원래 임베딩 크기로 축소
        self.fc2 = nn.Linear(ff_hidden_size, embed_size)
        # 활성화 함수로 ReLU 사용
        self.relu = nn.ReLU()
        
    def forward(self, x):
        # 입력 x에 대해 ReLU를 거친 뒤 다시 축소
        return self.fc2(self.relu(self.fc1(x)))

# 4. Transformer Block
class TransformerBlock(nn.Module):
    def __init__(self, embed_size, num_heads, ff_hidden_size, dropout):
        super().__init__()
        # 멀티헤드 어텐션 레이어
        self.attention = MultiHeadAttention(embed_size, num_heads)
        # Layer Normalization
        self.norm1 = nn.LayerNorm(embed_size)
        self.norm2 = nn.LayerNorm(embed_size)
        # Feedforward Network
        self.feed_forward = FeedForward(embed_size, ff_hidden_size)
        # 드롭아웃: 과적합 방지
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, value, key, query, mask):
        # 어텐션 + 잔차 연결 + LayerNorm
        attention = self.attention(query, key, value, mask)
        x = self.norm1(attention + query)
        # Feedforward Network + 잔차 연결 + LayerNorm
        forward = self.feed_forward(x)
        out = self.norm2(forward + x)
        return out

# 5. Positional Encoding
class PositionalEncoding(nn.Module):
    def __init__(self, embed_size, max_len=100):
        super().__init__()
        # 최대 길이와 임베딩 크기에 대한 포지셔널 인코딩 초기화
        self.encoding = torch.zeros(max_len, embed_size)
        self.encoding.requires_grad = False  # 학습되지 않도록 설정

        # 포지션 행렬 생성
        position = torch.arange(0, max_len).unsqueeze(1).float()
        # 주기를 계산하는 항
        div_term = torch.exp(torch.arange(0, embed_size, 2).float() * -(math.log(10000.0) / embed_size))
        
        # 짝수 인덱스에 대해 사인 적용
        self.encoding[:, 0::2] = torch.sin(position * div_term)
        # 홀수 인덱스에 대해 코사인 적용
        self.encoding[:, 1::2] = torch.cos(position * div_term)
        self.encoding = self.encoding.unsqueeze(0)  # 배치 차원 추가

    def forward(self, x):
        # 입력에 포지셔널 인코딩을 추가
        return x + self.encoding[:, :x.size(1), :].to(x.device)

# 6. Transformer Model
class Transformer(nn.Module):
    def __init__(self, embed_size, num_heads, ff_hidden_size, num_layers, vocab_size, max_len, dropout):
        super().__init__()
        # 임베딩 레이어
        self.embedding = nn.Embedding(vocab_size, embed_size)
        # 포지셔널 인코딩
        self.positional_encoding = PositionalEncoding(embed_size, max_len)
        # Transformer 블록을 여러 층으로 구성
        self.layers = nn.ModuleList(
            [TransformerBlock(embed_size, num_heads, ff_hidden_size, dropout) for _ in range(num_layers)]
        )
        # 최종 출력 레이어
        self.fc_out = nn.Linear(embed_size, vocab_size)
        # 드롭아웃
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, mask):
        # 임베딩 + 포지셔널 인코딩 적용
        out = self.embedding(x)
        out = self.positional_encoding(out)
        # 각 Transformer 블록 통과
        for layer in self.layers:
            out = layer(out, out, out, mask)
        # 최종 결과 반환
        return self.fc_out(out)

# 7. 모델 초기화
embed_size = 512  # 임베딩 차원
num_heads = 8  # 멀티헤드 어텐션의 헤드 수
ff_hidden_size = 2048  # 피드포워드 네트워크의 히든 레이어 크기
num_layers = 6  # Transformer 블록 수
vocab_size = 30522  # 어휘 크기
max_len = 100  # 문장의 최대 길이
dropout = 0.1  # 드롭아웃 비율

model = Transformer(embed_size, num_heads, ff_hidden_size, num_layers, vocab_size, max_len, dropout)

# 샘플 입력 데이터
sample_input = torch.randint(0, vocab_size, (2, 20))  # 배치 크기: 2, 문장 길이: 20
output = model(sample_input, mask=None)
print(output.shape)  # 출력 크기: (2, 20, vocab_size)