Transformer Pytorch 코드 리뷰

기본 Transformer 구조

 Pytorch에는 Transformer가 구현되어 있어, 이를 활용한 기본 구조를 만들면 다음과 같다.

# nn모듈을 활용한 Transformer 기본구조
import torch
import torch.nn as nn
import torch.optim as optim

class Transformer(nn.Module):
    def __init__(self, num_tokens, dim_model, num_heads, num_encoder_layers, num_decoder_layers, dropout_p, ):
        super().__init__()

        # Layers
        self.transformer = nn.Transformer(
            d_model=dim_model,
            nhead=num_heads,
            num_encoder_layers=num_encoder_layers,
            num_decoder_layers=num_decoder_layers,
            dropout=dropout_p,
        )

    def forward(self):
        pass
        
# dim_model - encoder와 decoder에서 정해진 입력과 출력의 크기 (default = 512)
# num_heads - Multi-head attention의 head 수이다.(default = 512)
# num_encdoer_layers - encoder의 층 개수
# num_decoder_layers - decoder의 층 개수
# dim_feedforward - feedforward network의 은닉층의 크기(default = 2048)

 nn.Transformer는 Transformer의 전체 구조에서 Multi-head attention과 Feed-Forward 계층을 처리함으로 Linear, Positional Encoding, Embedding은 따로 만들어야한다.

Positional Encoding

 Transformer는 입력 시퀀스의 순서를 신경쓰지 않음으로 순서 정보를 주입할 필요가 있다. 이를 위해 Positional Encoding을 하게 되는데, 위를 위해 sin과 cos함수를 활용한다.

sin, cos 함수를 사용하는 이유

• 이 함수들의 출력값은 입력값에 따라 달라지기에 출력값으로 입력값의 상대적인 위치를 알 수 있다.
• 이 함수들은 규칙적으로 증감하기에 딥러닝 모델이 이 규칙을 사용해 입력값의 상대적 위치를 쉽게 이용 가능하다.
• 출력값은 -1~1 범위이기 때문에 무한대 길이의 입력값도 출력 가능하다.

class PositionalEncoding(nn.Module):
    def __init__(self, dim_model, dropout_p, max_len):
        super().__init__()
        # 드롭 아웃
        self.dropout = nn.Dropout(dropout_p)

        # Encoding - From formula
        pos_encoding = torch.zeros(max_len, dim_model)
        positions_list = torch.arange(0, max_len, dtype=torch.float).view(-1, 1) # 0, 1, 2, 3, 4, 5
        division_term = torch.exp(torch.arange(0, dim_model, 2).float() * (-math.log(10000.0)) / dim_model) # 1000^(2i/dim_model)

        pos_encoding[:, 0::2] = torch.sin(positions_list * division_term)
        pos_encoding[:, 1::2] = torch.cos(positions_list * division_term)

        # Saving buffer (same as parameter without gradients needed)
        pos_encoding = pos_encoding.unsqueeze(0).transpose(0, 1)
        self.register_buffer("pos_encoding",pos_encoding)

    def forward(self, token_embedding: torch.tensor) -> torch.tensor:
        # Residual connection + pos encoding
        return self.dropout(token_embedding + self.pos_encoding[:token_embedding.size(0), :])

 

• positions_list는 수식의 분자인 p에 해당하는 단어의 위치(순서)를 의미하는 변수
• division_term의 dim_model(default=512)은 모델의 입출력 크기로, 수식에서 d에 해당
• divistion_term에서 0~dim_model까지의 짝수들을 만들어내는데, i가 홀수일 경우 짝수로 변경해 사용하기 때문에 여기서 짝수만 만든다.
• pos_encoding에서 짝수번째 벡터는 sin, 홀수번째 벡터는 cos을 통해 인코딩한다. (dim_model이 시퀀스의 길이가 아니라 임베딩된 단어의 차원이므로, 각 단어별로 생성된 임베딩 벡터에서 짝수번째와 홀수번째 단어에 따라 함수가 사용된다.)
• pos_encoding에 unsqeeze와 transpose를 사용해 차원을 [max_len, dim_model] 에서 [max_len, 1, dim_model]
• 로 변환한다.
• register_buffer로 layer를 등록하면 Optimizer 업데이트 없이 하나의 layer로 작용한다. 또한 state_dict()로 확인 가능하고, GPU 연산이 가능하다.\nn.Module.register_buffer('attribute_name', tensor)
  • nn.Module.register_buffer('attribute_name', tensor)
  • 모듈 내에서 tensor는 self.attribute_name으로 접근 가능하다
  • tensor는 학습되지 않는다
  • model.cuda() 시에 tensor도 GPU로 연산된다.

따라서 이 class에 입력된 max_len과 dim_model 파라미터에 의해 forward에 입력될 token보다 크거나 같은 positional encoding vector를 만들고, forward 시에 token길이에 맞추어 slice 이후 결합하는 구조이다.

 

Transformer 완성

class Transformer(nn.Module):
    # Constructor
    def __init__( self, num_tokens, dim_model, num_heads, num_encoder_layers, num_decoder_layers, dropout_p, ):
        super().__init__()

        # INFO
        self.model_type = "Transformer"
        self.dim_model = dim_model

        # LAYERS
        self.positional_encoder = PositionalEncoding(dim_model=dim_model, dropout_p=dropout_p, max_len=5000)
        self.embedding = nn.Embedding(num_tokens, dim_model)
        self.transformer = nn.Transformer(
            d_model=dim_model,
            nhead=num_heads,
            num_encoder_layers=num_encoder_layers,
            num_decoder_layers=num_decoder_layers,
            dropout=dropout_p,
        )
        self.out = nn.Linear(dim_model, num_tokens)

    def forward(self, src, tgt, tgt_mask=None, src_pad_mask=None, tgt_pad_mask=None):
        # Src, Tgt size 는 반드시 (batch_size, src sequence length) 여야 합니다.

        # Embedding + positional encoding - Out size = (batch_size, sequence length, dim_model)
        src = self.embedding(src) * math.sqrt(self.dim_model)
        tgt = self.embedding(tgt) * math.sqrt(self.dim_model)
        src = self.positional_encoder(src)
        tgt = self.positional_encoder(tgt)

        src = src.permute(1,0,2)
        tgt = tgt.permute(1,0,2)

        # Transformer blocks - Out size = (sequence length, batch_size, num_tokens)
        transformer_out = self.transformer(src, tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_pad_mask, tgt_key_padding_mask=tgt_pad_mask)
        out = self.out(transformer_out)

        return out

 

• layer는 크게 Positional Encoder, Embedding, Transformer로 나뉘어져있다.
• forward 시에 poistional encoding을 하기 전 token에 √𝑑를 곱하는데(scaling), 아래 링크에 따르면 embedding 시 벡터값들은 평균이 0, 분산이 √𝑑로 표준화되는데, -1~1사이의 값을 갖는 positional encoding 벡터와의 스케일을 맞추기 위해 곱하는 것으로 이해되었다.https://datascience.stackexchange.com/questions/87906/transformer-model-why-are-word-embeddings-scaled-before-adding-positional-encod/87909#87909

Transformer model: Why are word embeddings scaled before adding positional encodings?