공일 2025.3.18 (Transformer 구현)

공부 목적으로 Transformer를 구현하고 한-영 번역기를 만들어 봤다.

막상 Transformer구현 자체는 그리 어렵지 않았고, 데이터 처리를 구현하면서 꽤 애먹었다.

 

Transformer Code

import torch as th
import math

_SQRT_EPS = 1e-5


class LayerNorm(th.nn.Module) :
    def __init__(self, n_dim:int, eps:float=_SQRT_EPS) :
        super().__init__()
        self.eps = eps

        self.gamma = th.nn.Parameter(th.ones(size=(n_dim,), dtype=th.float32), requires_grad=True)
        self.beta = th.nn.Parameter(th.zeros(size=(n_dim,), dtype=th.float32), requires_grad=True)

    def forward(self, x:th.Tensor) :
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, correction=0)

        x = (x - mean) / th.sqrt(var + self.eps)
        x = x * self.gamma + self.beta

        return x


class MultiHeadAttention(th.nn.Module) :
    def __init__(
            self,
            n_head:int,
            n_input:int,
            n_key:int,
            n_value:int,
            n_output:int,
        ) :
        super().__init__()

        self.n_head = n_head
        self.n_input = n_input
        self.n_key = n_key
        self.n_value = n_value
        self.n_output = n_output

        self._dot_scaler = 1 / math.sqrt(self.n_key)

        self.w_query    = th.nn.Linear(self.n_input, self.n_head * self.n_key, bias=False)
        self.w_key      = th.nn.Linear(self.n_input, self.n_head * self.n_key, bias=False)
        self.w_value    = th.nn.Linear(self.n_input, self.n_head * self.n_value, bias=False)
        self.project_output = th.nn.Linear(self.n_head * self.n_value, self.n_output)

    def _to_batched_shape(self, x:th.Tensor, transpose:bool=False) :
        '''
        x: (n_batch, n_seq, n_head, n_dim)
        '''
        _, n_seq, _, n_dim = x.shape
        x = x.transpose(1, 2)
        # (n_batch, n_head, n_seq, n_dim)
        if not transpose :
            x = x.contiguous()
            x = x.view(-1, n_seq, n_dim)
            # (*, n_seq, n_dim)
        else :
            x = x.transpose(2, 3)
            x = x.contiguous()
            x = x.view(-1, n_dim, n_seq)
            # (*, n_dim, n_seq)
        return x
    
    def forward(
            self,
            q:th.Tensor,
            k:th.Tensor,
            v:th.Tensor,
            mask:th.Tensor,
        ) :
        '''
        q: (n_batch, n_seq, n_input)
        k: (n_batch, n_seq, n_input)
        v: (n_batch, n_seq, n_input)
        mask: bool(n_batch, n_seq, n_seq)
        '''
        n_batch, n_seq, _ = q.shape
        def decompose(a:th.Tensor) :
            '''
            a: (n_batch, n_seq, n_head * n_dim)
            '''
            return a.view(n_batch, n_seq, self.n_head, -1)

        # shape tensor
        q   = self._to_batched_shape(decompose(self.w_query(q)))
        k_T = self._to_batched_shape(decompose(self.w_key(k)), transpose=True)
        v   = self._to_batched_shape(decompose(self.w_value(v)))
        mask = mask.unsqueeze(1)
        # (n_batch, 1, n_seq, n_seq)
        mask = mask.expand(n_batch, self.n_head, n_seq, n_seq)
        mask = mask.contiguous().view(-1, n_seq, n_seq)
        # (*, n_seq, n_seq)

        # compute weighted sum
        # scaled_softmax(q * k^T) * v
        weight = th.bmm(q, k_T) * self._dot_scaler
        weight[mask] = -th.inf
        weight = th.softmax(weight, dim=-1)
        y = th.bmm(weight, v)
        # (-1, n_seq, n_value)

        # unshape tensor
        y = y.view(-1, self.n_head, n_seq, self.n_value)
        y = y.transpose(1, 2).contiguous()
        # (n_batch, n_seq, n_head, n_value)
        y = y.view(-1, n_seq, self.n_head * self.n_value)
        # (n_batch, n_seq, n_head * n_value)

        y = self.project_output(y)
        # (n_batch, n_seq, n_output)
        return y


class EncoderStack(th.nn.Module) :
    def __init__(
        self,
        n_head:int,
        n_model:int,
        n_ff:int,
        n_ff_dim:int,
        ff_activ_fn:th.nn.Module,
        drop_prob:float,
    ) :
        super().__init__()
        self.n_head = n_head
        self.n_model = n_model
        self.drop_prob = drop_prob

        if self.n_model % self.n_head != 0 :
            raise Exception()
        
        self.self_attention = MultiHeadAttention(
            n_head=self.n_head,
            n_input=self.n_model,
            n_key=self.n_model//self.n_head,
            n_value=self.n_model//self.n_head,
            n_output=self.n_model,
        )
        
        ff_arch = [self.n_model] + [n_ff_dim for _ in range(n_ff-1)] + [self.n_model]
        self.ff_layers = []
        for idx in range(n_ff) :
            self.ff_layers.append(th.nn.Linear(ff_arch[idx], ff_arch[idx + 1]))
            self.ff_layers.append(ff_activ_fn())
        self.ff_layers.pop()
        self.ff_layers = th.nn.ModuleList(self.ff_layers)

        self.layer_norm1 = LayerNorm(self.n_model)
        self.layer_norm2 = LayerNorm(self.n_model)
    
    def forward(self, x:th.Tensor, padding_mask:th.Tensor) :
        '''
        x: (n_batch, n_seq, n_model)
        padding_mask: bool(n_batch, n_seq, n_seq)
        '''

        # self attention
        x_src = x
        x = self.self_a

 

번역기 학습 결과

이 책은 음식을 요리하는 방법과 그것을 어떻게 맛있게 먹을지를 포함합니다.
**** Input ****
이 책은 음식을 요리하는 방법과 그것을 어떻게 맛있게 먹을지를 포함합니다.
**** Output ****
This book includes how to cook the food and how to enjoy it.

지난 주말에는 가족과 함께 산을 오르며 아름다운 풍경을 감상했고, 그 경험은 잊을 수 없는 추억이 되었습니다.
**** Input ****
지난 주말에는 가족과 함께 산을 오르며 아름다운 풍경을 감상했고, 그 경험은 잊을 수 없는 추억이 되었습니다.
**** Output ****
Last weekend, I enjoyed the mountain with my family and enjoyed the mountains, and it became a memory that I can't forget the most.

그는 내일 파티에 올 것이라고 약속했습니다.
**** Input ****
그는 내일 파티에 올 것이라고 약속했습니다.
**** Output ****
He promised to come to the party tomorrow.

모든 사람이 평화를 누리기를 바랍니다.
**** Input ****
모든 사람이 평화를 누리기를 바랍니다.
**** Output ****
I hope everyone enjoys peace.

어려운 문제에 직면했을 때, 서로 도와가며 해결책을 찾는 협력의 중요성을 다시 한 번 깨달았습니 다.
**** Input ****
어려운 문제에 직면했을 때, 서로 도와가며 해결책을 찾는 협력의 중요성을 다시 한 번 깨달았습니 다.
**** Output ****
When we faced a difficult question, we realized the importance of finding a solution and finding a solution.

 

 

학습 목적으로 빠르게 만들어 본거다 보니 번역이 정확하지는 않지만, 최소한의 번역이 이루어 지는것은 확인할 수 있다.

 

이후 계획

transformer를 공부했으니 RL 쪽에서 transformer를 이용하는 연구 (e.g. Decision Transformer)들을 몇개 찾아볼 예정이다.

더불어 vit 같은 것도 한번 찾아 볼 예정이다.