728x90

모델을 돌리기 위해서 만든 건 아니고, 이전에 공부한 내용을 코드로 구현해보면서 이해를 깊이하는 시간을 갖고자 만들어보았다.
따라서 실제 데이터 넣고 학습했을 때 구데기일 수 있다는 점을 참고해주시길...

1. Attention 구현하기

$$Attention = softmax(\frac{QK^{T}}{\sqrt{ d_{k} }})V$$
이 공식에 맞게 구현하면 된다. 여기서 $d_{k}$ 는 k의 차원수이다.
차례차례 보면, 순서대로 차근차근 진행하면 된다. 마스크 적용에 관한 건은 뒤에서.

def scaled_dot_product_attention(query, key, value, mask=None):
    # QK^T
    matmul_qk = tf.matmul(query, key, transpose_b=True) # 두 개 중 뒤에 해당하는 값 transpose

    # QK^T / sqrt(d_k)
    d_k = tf.cast(tf.shape(key)[-1], tf.float32)
    scaled_attention_logits = matmul_qk / tf.math.sqrt(d_k)

    # softmax
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)

    # multiply with V
    output = tf.matmul(attention_weights, value)

    return output, attention_value

2. Multi head Attention 구현하기

Multi head attention은 위에서 구현한 attention을 호출해서 self-attention을 여러 헤드가 독자적으로 계산하도록 하는 구현 방법이다. 이 클래스 자체를 하나의 레이어로 볼 것이므로 tf.keras.layers.Layer을 상속받아 구현한다.

class MultiHeadAttention(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.d_model = d_model

        # d_model % num_heads 는 언제나 0이어야 함! 평등하게 나눠서 계산할 것
        assert self.d_model % self.num_heads == 0, "d_model must be divisiable by num_heads"

        self.depth = self.d_model // self.num_heads

        # q, k, v를 업데이트할 가중치
        self.wq = tf.keras.layers.Dense(d_model)
        self.wk = tf.keras.layers.Dense(d_model)
        self.wv = tf.keras.layers.Dense(d_model)

        self.dense = tf.keras.layers.Dense(d_model)

    def split_heads(self, x, batch_size):
        # d_model을 head 개수만큼 나눠서 입력을 multi head로 만들 초석을 다짐
        # shape = (64, x_length, num_heads, d_model // num_heads)
        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
        return tf.transpose(x, perm=[0, 2, 1, 3])

    # multi-head-attention layer 호출
    def call(self, query, key, value, mask=None):
        batch_size = tf.shape(query)[0]

        # Q, K, T 구함
        query = self.wq(query)
        key = self.wk(key)
        value = self.wv(value)

        # multi head에 맞게 나누고 reshape
        query = self.split_heads(query, batch_size)
        key = self.split_heads(key, batch_size)
        value = self.split_heads(value, batch_size)

        # Attention score 계산
        scaled_attention, attention_weights = scaled_dot_product_attention(query, key, value, mask)

        # 계산 쉽게 하기 위해 바꿨던 위치 복원
        scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])
        concat_attention = tf.reshape(scaled_attention, (batch_size, -1, self.d_model))
        output = self.dense(concat_attention)

        return output, attention_weights

3. Feed Forward 구현하기

Encoder와 Decoder 마지막 단계에서 Feed Forward layer가 존재한다.

class PositionwiseFeedForward(tf.keras.layers.Layer):
    def __init__(self, d_model, d_ffn=2048):
        super(PositionwiseFeedForward, self).__init__()
        self.d_model = d_model
        self.d_ffn = d_ffn

        self.dense1 = layers.Dense(self.d_ffn, activation="relu")
        self.dense2 = layers.Dense(self.d_model)

    def call(self, x):
        x = self.dense1(x)
        x = self.dense2(x)
        return x

Feed forward 네트워크는 단순하게 두 개의 fully connected layer를 연결하는 방식으로 구성하였다.

4. Encoder 구현하기

4.1. Encoder layer

구현해야할 Encoder Layer를 보면, Multi-head Attention과 normalization 레이어가 있고 feed forward layer로 이어진다. 그리고 residual connection을 이용해서 layer의 입력 데이터를 출력 데이터에 반영하는 식으로 구성되어 있다.

위 단계에서 multi-head connection과 feed forward를 구형하두었으므로 이를 코드로 표현하면 아래와 같다.

class EncoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, d_ffn, dropout_rate=0.1):
        super(EncoderLayer, self).__init__()
        # 사용할 layer 선언
        # Multi-Head Attention
        self.multi_head_attention = MultiHeadAttention(d_model, num_heads)

        # FeedForward Network
        self.feedforward = PositionwiseFeedForward(d_model, d_ffn)

        # Dropout
        self.dropout1 = layers.Dropout(dropout_rate)
        self.dropout2 = layers.Dropout(dropout_rate)

        # Layer Normalization
        self.layer_norm1 = layers.LayerNormalization(epsilon=1e-6)
        self.layer_norm2 = layers.LayerNormalization(epsilon=1e-6)

    def call(self, x, mask, training):
        # multi-head attention 실행
        attn_output, _ = self.multi_head_attention(x, x, x, mask)
        # dropout
        attn_output = self.dropout1(attn_output, training=training)
        # 입력 x와 attention output을 더해서 residual connection 구현 + norm
        out1 = self.layer_norm1(x + attn_output)

        # feed forward network
        ffn_output = self.feedforward(out1)
        # dropout
        ffn_output = self.dropout2(ffn_output, training=training)
        # attention 출력 out1과 feead forward output을 더해서 residual connection 구현 + norm
        out2 = self.layer_norm2(out1 + ffn_output)

        return out3

이제 이 encoder layer를 여러 개 붙여서 encoder를 구성한다.

먼저 encoder 입력 벡터를 embeding 해야한다. 입력 벡터를 고정 크기의 dense 벡터로 변환하는 과정이다.
그리고 positional encoding의 수행결과와 embedding을 먼저 합쳐주고, 이를 encoder layer들의 입력으로 사용하도록 한다.

class Encoder(tf.keras.layers.Layer):
    def __init__(
        self, num_layers, d_model, num_heads, d_ffn, vocab_size, dropout_rate=0.1
    ):
        super(Encoder, self).__init__()

        # 입력 임베딩
        self.embedding = tf.keras.layers.Embedding(
            input_dim=vocab_size, output_dim=d_model
        )  # 단어집의 크기가 10000일 경우

        # Positional Encoding
        self.pos_encoding = self.add_weight(
            "pos_encoding", shape=[1, 10000, d_model], trainable=False
        )

        self.encoder_layers = [
            EncoderLayer(d_model, num_heads, d_ffn, dropout_rate)
            for _ in range(num_layers)
        ]

        self.dropout = tf.keras.layers.Dropout(dropout_rate)  

    def call(self, x, mask, training):
        # embedding + positional encoding
        x = self.embedding(x) + self.pos_encoding[:, : tf.shape(x)[1], :]
        x = self.dropout(x, training=training)
        for layer in self.encoder_layers:
            x = layer(x, mask, training)
        return x

4.2 Encoder Padding Mask

def create_padding_mask(seq):
    mask = tf.cast(tf.math.equal(seq, 0), tf.float32)
    return mask[:, tf.newaxis, tf.newaxis, :]  # shape = (batch_size, 1, 1, seq_len)

Sequence 길이를 통일하기 위해서 만들어진 padding이다. Transformer는 고정길이 sequence를 입출력으로 두기 때문에 시퀀스 길이가 들쭉날쭉해서는 안된다.

5. Decoder 구현하기

5.1. Decoder Layer

Decoder layer는 두 개의 multi-head attention과 하나의 feed forward network로 구성되어 있다. 이때, decoder layer의 첫 번째 multi-head attention은 masked multi-head attention 이라고 하는데, decoder의 입력이 target 시퀀스인 것을 고려하였을 때, 현재 모델이 추론해야하는 순서의 단어보다 미래의 정보를 갖지 못하도록 가리는 역할을 한다. 미래의 정보를 알면 이미 답을 아는 것과 다름없기 때문이다.

class DecoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, d_ffn, dropout_rate=0.1):
        super(DecoderLayer, self).__init__()

        # Masked Multi-head Attention
        self.masked_multi_head_attention = MultiHeadAttention(
            d_model=d_model, num_heads=num_heads
        )

        # Multi-Head Attention
        self.multi_head_attention = MultiHeadAttention(
            d_model=d_model, num_heads=num_heads
        )

        # FeedForward Network
        self.feedforward = PositionwiseFeedForward(d_model=d_model, d_ffn=d_ffn)

        # Dropout
        self.dropout1 = layers.Dropout(dropout_rate)
        self.dropout2 = layers.Dropout(dropout_rate)
        self.dropout3 = layers.Dropout(dropout_rate)

        # Layer Normalization
        self.layer_norm1 = layers.LayerNormalization(epsilon=1e-6)
        self.layer_norm2 = layers.LayerNormalization(epsilon=1e-6)
        self.layer_norm3 = layers.LayerNormalization(epsilon=1e-6)

    def call(self, x, encoder_output, look_ahead_mask, padding_mask, training):
        # Masked Multi-Head Attention
        attn1, _ = self.masked_multi_head_attention(x, x, x, mask=look_ahead_mask)
        attn1 = self.dropout1(attn1, training=training)
        out1 = self.layer_norm1(x + attn1)

        attn2, _ = self.multi_head_attention(
            out1, encoder_output, encoder_output, mask=padding_mask
        )
        attn2 = self.dropout2(attn2, training=training)
        out2 = self.layer_norm2(out1 + attn2)

        ffn_output = self.feedforward(out2)
        ffn_output = self.dropout3(ffn_output, training=training)
        out3 = self.layer_norm3(out2 + ffn_output)

        return out3

5.2. Decoder padding mask

def create_padding_mask(seq):
    mask = tf.cast(tf.math.equal(seq, 0), tf.float32)
    return mask[:, tf.newaxis, tf.newaxis, :]  # shape = (batch_size, 1, 1, seq_len)

이 부분은 Encoder padding mask와 같다. 다만 decoder sequence는 encoder sequence보다 길이가 하나 짧다는 사실을 염두에 두어야한다.

5.3. Look-ahead mask

def create_look_ahead_mask(size):
    """
    tf.ones((4,4)) -> [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]
    tf.linalg.band_part,
      num_lower = -1 : 대각선 아래 모든 요소를 남김
      num_upper = 0 : 대각선 위 모든 요소를 0으로 처리리 -> [[1, 0, 0, 0], [1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 1]]
    1 - tf.linalg.band_part(~) : 에서 빼줘서 값 반전 -> [[0, 1, 1, 1], [0, 0, 1, 1], [0, 0, 0, 1], [0, 0, 0, 0]]
    1로 마스킹됨 -> 나중에 1e-6 이랑 곱해져서 참조 못하게 될 예정
    """
    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), num_lower=-1, num_upper=0)
    return mask

look-ahead 마스크는 학습 중에 decoder가 추론해야할 값보다 더 미래의 정보들을 보지 못하도록 막는 역할을 한다. decoder에서 첫 번째 multi-head attention을 수행할 때 미래의 값을 보지 못하도록, 정확히는 미래의 정보값이 영향을 미치지 못하도록 아주 작게 만들어 학습의 영향권에서 빼는 역할을 한다.

6. Transformer

class Transformer(tf.keras.Model):
    def __init__(
        self,
        num_enc_layers,
        num_dec_layers,
        d_model,
        num_heads,
        d_ffn,
        input_vocab_size,
        target_vocab_size,
        max_pos_enc,
        dropout_rate=0.1,
    ):
        super(Transformer, self).__init__()

        # Encoder
        self.encoder = Encoder(
            num_enc_layers, d_model, num_heads, d_ffn, input_vocab_size, dropout_rate
        )

        # Decoder
        self.decoder = Decoder(
            num_dec_layers, d_model, num_heads, d_ffn, target_vocab_size, dropout_rate
        )

        # output
        self.dense = layers.Dense(target_vocab_size)

    def call(
        self,
        enc_input,
        dec_input,
        encoder_padding_mask,
        look_ahead_mask,
        decoder_padding_mask,
        training,
    ):
        # Encoder
        encoder_output = self.encoder(enc_input, encoder_padding_mask, training)

        # Decoder
        decoder_output = self.decoder(
            dec_input, encoder_output, look_ahead_mask, decoder_padding_mask, training
        )

        # output
        output = self.dense(decoder_output)

        return output

이제껏 구현한 레이어들을 순차적으로 쌓아주기만 하면 된다.

사실 만든 모델을 딱히 테스트(...) 하지는 않았고 구조를 코드로 구현해보는데에 의의를 뒀다.
따로 데이터를 구해서 학습해보게 된다면 그것도 추가하겠다.

728x90

저작자표시 비영리 변경금지

'🐬 ML & Data > 📘 논문 & 모델 리뷰' 카테고리의 다른 글

[Paper Review] Transformer - Attention is All You Need (1)	2024.12.30
[Paper Review] Mamba - Linear Time Sequence Modeling with Selective State Spaces 2 (1)	2024.12.11
[Paper Review] Mamba - Linear Time Sequence Modeling with Selective State Spaces 1 (1)	2024.12.11
[Paper Review] Transforming Cooling Optimization for Green Data Center via Deep Reinforcement Learning (0)	2023.08.07
[Model Review] TadGAN(Time series Anomaly Detection GAN) (0)	2023.05.17

[논문 구현] Transformer 텐서플로우로 구현하기

1. Attention 구현하기

2. Multi head Attention 구현하기

3. Feed Forward 구현하기

4. Encoder 구현하기

4.1. Encoder layer

4.2 Encoder Padding Mask

5. Decoder 구현하기

5.1. Decoder Layer

5.2. Decoder padding mask

5.3. Look-ahead mask

6. Transformer

'🐬 ML & Data > 📘 논문 & 모델 리뷰' 카테고리의 다른 글

티스토리툴바