본문 바로가기
머신러닝(Machine Learning)

NLP 트랜스포머 코드 스터디 리뷰 (1) transformer.py

by Blaze_블즈 2023. 8. 18.

안녕하세요

블레이즈 테크노트 

블레이즈 입니다. 

 

트랜스포머 논문에 대해서는,

제가 지난 여러 포스팅에서 설명했습니다. 

 

어텐션메커니즘1

어텐션메커니즘2

트랜스포머_포지셔널인코딩

트랜스포머_셀프어텐션

트랜스포머_멀티헤드어텐션

 

다시 짚고 넘어가자면 NLP 자연어 처리에서 아주 혁신적인 개념이었죠. 

 

Attention is all you need 가 구글에서 발표한 논문인만큼 

이 논문에서 사용된 코드가 tensorflow 공식으로 등록되어 있습니다. 

 

아래는 그 코드의 내용입니다. 

https://github.com/tensorflow/tensor2tensor

 

GitHub - tensorflow/tensor2tensor: Library of deep learning models and datasets designed to make deep learning more accessible a

Library of deep learning models and datasets designed to make deep learning more accessible and accelerate ML research. - GitHub - tensorflow/tensor2tensor: Library of deep learning models and data...

github.com

 

이 사이트를 보면 파일이 아주 많습니다. 

 

오늘은 그 중에서도 모델의 기본 골자를 구성하는 transformer.py 를 함께 살펴보겠습니다. 

https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/models/transformer.py

 

 

# coding=utf-8
# Copyright 2023 The Tensor2Tensor Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Transformer model from "Attention Is All You Need".

The Transformer model consists of an encoder and a decoder. Both are stacks
of self-attention layers followed by feed-forward layers. This model yields
good results on a number of problems, especially in NLP and machine translation.

See "Attention Is All You Need" (https://arxiv.org/abs/1706.03762) for the full
description of the model and the results obtained with its early version.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from six.moves import range  # pylint: disable=redefined-builtin

from tensor2tensor.data_generators import librispeech
from tensor2tensor.layers import common_attention
from tensor2tensor.layers import common_hparams
from tensor2tensor.layers import common_layers
from tensor2tensor.layers import modalities
from tensor2tensor.layers import transformer_layers
from tensor2tensor.layers import transformer_memory
from tensor2tensor.utils import beam_search
from tensor2tensor.utils import expert_utils
from tensor2tensor.utils import mlperf_log
from tensor2tensor.utils import registry
from tensor2tensor.utils import t2t_model

import tensorflow.compat.v1 as tf
from tensorflow.compat.v1 import estimator as tf_estimator

# pylint: disable=g-direct-tensorflow-import
from tensorflow.python.ops import inplace_ops
from tensorflow.python.util import nest
# pylint: enable=g-direct-tensorflow-import

먼저 필요한 모듈을 import합니다. 

이 논문을 보시면 굉장히 많은 모듈이 있습니다. 

 

1.  common_attention 모듈을 임포트 해줍니다. 이 모듈에 대한 설명은 여기에 설명하겠습니다. 

2. common_layers 모듈을 임포트 해줍니다. 이 모듈에 대한 설명은 여기에 설명하겠습니다. 

3. transformer_layers 모듈을 임포트 해줍니다. 이 모듈에 대한 설명은 여기에 설명하겠습니다. 

 

이 외의 나머지 모듈에 대해선 아래의 코드를 보면서 직접 사용될 때 설명하도록 하겠습니다. 

 

# Alias some commonly reused layers, here and elsewhere.
transformer_prepare_encoder = transformer_layers.transformer_prepare_encoder
transformer_encoder = transformer_layers.transformer_encoder
transformer_ffn_layer = transformer_layers.transformer_ffn_layer

 

이 코드에서는 transformer_prepare_encoder를 정의합니다. 

이것은 결국 transformer_layers 모듈에 정의되어 있던 transformer_prepare_encoder 함수를 그대로 가져옵니다. 

이 함수의 역할은 간단히 말하면 인코더에 들어가기 전에 전처리를 해주는 거라고 볼 수 있습니다. 

 

다음으로는 transformer_encoder를 정의해주는데

이 역시 transformer_layers 모듈에 정의되어 있던 transformer_encoder 함수를 사용합니다. 

 

마지막으로 transformer_layers.transformer_ffn_layer 에서 transformer_ffn_layer를 가져옵니다. 

 

이것으로 인코더에 들어갈 layer들은 모두 준비가 된 것 같습니다. 

def transformer_encode(encoder_function, inputs, target_space, hparams,
                       attention_weights=None, features=None, losses=None,
                       prepare_encoder_fn=None, **kwargs):
  """Encode transformer inputs.

  Args:
    encoder_function: the encoder function
    inputs: Transformer inputs [batch_size, input_length, 1, hidden_dim] which
      will be flattened along the two spatial dimensions.
    target_space: scalar, target space ID.
    hparams: hyperparameters for model.
    attention_weights: weight to store attention to.
    features: optionally pass the entire features dictionary as well. This is
      needed now for "packed" datasets.
    losses: optional list onto which to append extra training losses
    prepare_encoder_fn: optional, alternative to transformer_prepare_encoder.
    **kwargs: additional arguments to pass to encoder_function

  Returns:
    Tuple of:
        encoder_output: Encoder representation.
            [batch_size, input_length, hidden_dim]
        encoder_decoder_attention_bias: Bias and mask weights for
            encoder-decoder attention. [batch_size, input_length]
  """
  inputs = common_layers.flatten4d3d(inputs)

  if not prepare_encoder_fn:
    prepare_encoder_fn = transformer_prepare_encoder
  encoder_input, self_attention_bias, encoder_decoder_attention_bias = (
      prepare_encoder_fn(
          inputs, target_space, hparams, features=features))

  mlperf_log.transformer_print(
      key=mlperf_log.MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
      value=hparams.layer_prepostprocess_dropout,
      hparams=hparams)

  encoder_input = tf.nn.dropout(encoder_input,
                                1.0 - hparams.layer_prepostprocess_dropout)

  attn_bias_for_padding = None
  # Otherwise the encoder will just use encoder_self_attention_bias.
  if hparams.unidirectional_encoder:
    attn_bias_for_padding = encoder_decoder_attention_bias

  encoder_output = encoder_function(
      encoder_input,
      self_attention_bias,
      hparams,
      nonpadding=features_to_nonpadding(features, "inputs"),
      save_weights_to=attention_weights,
      make_image_summary=not common_layers.is_xla_compiled(),
      losses=losses,
      attn_bias_for_padding=attn_bias_for_padding,
      **kwargs)

  return encoder_output, encoder_decoder_attention_bias

 

지금부터 transformer_encode라는 함수의 정의가 시작됩니다. 

  1. 입력 매개변수:
    • encoder_function: 실제 인코딩을 수행하는 함수입니다 
    • inputs: 인코딩이 필요한 초기 입력 텐서입니다.
    • target_space: 대상 공간을 식별합니다 
    • hparams: 모델의 하이퍼파라미터입니다.
    • 기타 매개변수들은 주의 깊게 처리하거나 저장할 필요가 있는 추가 정보를 위해 사용됩니다.
  1.  
  2. 전처리:
    • 입력은 4D 텐서에서 3D 텐서로 재구성됩니다.  common_layers.flatten4d3d가 이를 수행하는 함수입니다. 
    • 인코더 입력을 준비하기 위한 함수가 제공되지 않은 경우 기본값으로 transformer_prepare_encoder를 사용합니다.
      • transformer_prepare_encoder 는 입력 데이터에 대한 전처리 역할을 하는데 주로 셀프 어텐션 바이어스 등을 계산합니다. 
     

 

3. 드롭 아웃 : 오버피팅을 방지하기 위해 인코더 입력에 dropout 함수 적용

 

4. encoder_function()으로 encoder_output 반환

def transformer_decode(decoder_function,
                       decoder_input,
                       encoder_output,
                       encoder_decoder_attention_bias,
                       decoder_self_attention_bias,
                       hparams,
                       attention_weights=None,
                       cache=None,
                       decode_loop_step=None,
                       nonpadding=None,
                       losses=None,
                       **kwargs):
  """Decode Transformer outputs from encoder representation.

  Args:
    decoder_function: the decoder function
    decoder_input: inputs to bottom of the model. [batch_size, decoder_length,
      hidden_dim]
    encoder_output: Encoder representation. [batch_size, input_length,
      hidden_dim]
    encoder_decoder_attention_bias: Bias and mask weights for encoder-decoder
      attention. [batch_size, input_length]
    decoder_self_attention_bias: Bias and mask weights for decoder
      self-attention. [batch_size, decoder_length]
    hparams: hyperparameters for model.
    attention_weights: weight to store attention to.
    cache: dict, containing tensors which are the results of previous
      attentions, used for fast decoding.
    decode_loop_step: An integer, step number of the decoding loop. Only used
      for inference on TPU.
    nonpadding: optional Tensor with shape [batch_size, decoder_length]
    losses: optional list onto which to append extra training losses
    **kwargs: additional arguments to pass to decoder_function

  Returns:
    Final decoder representation. [batch_size, decoder_length, hidden_dim]
  """
  mlperf_log.transformer_print(
      key=mlperf_log.MODEL_HP_LAYER_POSTPROCESS_DROPOUT,
      value=hparams.layer_prepostprocess_dropout,
      hparams=hparams)
  decoder_input = tf.nn.dropout(decoder_input,
                                1.0 - hparams.layer_prepostprocess_dropout)

  decoder_output = decoder_function(
      decoder_input,
      encoder_output,
      decoder_self_attention_bias,
      encoder_decoder_attention_bias,
      hparams,
      cache=cache,
      decode_loop_step=decode_loop_step,
      nonpadding=nonpadding,
      save_weights_to=attention_weights,
      losses=losses,
      **kwargs)

  if (common_layers.is_xla_compiled() and
      hparams.mode == tf_estimator.ModeKeys.TRAIN):
    # TPU does not react kindly to extra dimensions.
    # TODO(noam): remove this once TPU is more forgiving of extra dims.
    return decoder_output
  else:
    # Expand since t2t expects 4d tensors.
    return tf.expand_dims(decoder_output, axis=2)

다음으로는 transformer_decode 함수의 정의입니다. 

 

transformer_decode 함수는 transformer_encode와 유사합니다. 

파라미터만 약간 다른데 그 부분만 설명하도록 하겠습니다. 

 

transformer_decode 함수의 핵심적인 부분은 결국

decoder_function에 의해 decoder_output이 정의되고 이를 리턴한다는 것입니다. 

decoder_function은 transformer_decode의 파라미터로 주어집니다. 

그 외의 파라미터로는 아래와 같습니다. 

  • decoder_input: 모델의 하단에 대한 입력입니다.
  • encoder_output: 인코더에서 얻은 표현입니다.
  • encoder_decoder_attention_bias: 인코더-디코더 어텐션에 대한 바이어스 및 마스크 가중치입니다.
  • decoder_self_attention_bias: 디코더 셀프 어텐션에 대한 바이어스 및 마스크 가중치입니다.
  • hparams: 모델의 하이퍼파라미터입니다.
  • 기타 매개변수들은 추가 정보나 처리, 저장을 위해 사용됩니다.

 

나머지는 로그를 기록하고 드롭 아웃을 하는 부분인 것 같네요. 

지금부터는 Transformer 클래스 정의 입니다. 

@registry.register_model
class Transformer(t2t_model.T2TModel):
  """Attention net.  See file docstring."""

  def __init__(self, *args, **kwargs):
    super(Transformer, self).__init__(*args, **kwargs)
    self.attention_weights = {}  # For visualizing attention heads.
    self.recurrent_memory_by_layer = None  # Override to enable recurrent memory
    self._encoder_function = transformer_encoder
    self._decoder_function = transformer_decoder
    self._init_cache_fn = _init_transformer_cache
    self._prepare_encoder_fn = transformer_prepare_encoder
    self._prepare_decoder_fn = transformer_prepare_decoder

  def encode(self, inputs, target_space, hparams, features=None, losses=None):
    """Encode transformer inputs, see transformer_encode."""
    return transformer_encode(
        self._encoder_function, inputs, target_space, hparams,
        attention_weights=self.attention_weights,
        features=features, losses=losses,
        prepare_encoder_fn=self._prepare_encoder_fn)

  def decode(self,
             decoder_input,
             encoder_output,
             encoder_decoder_attention_bias,
             decoder_self_attention_bias,
             hparams,
             cache=None,
             decode_loop_step=None,
             nonpadding=None,
             losses=None,
             **kwargs):
    """Decode Transformer outputs, see transformer_decode."""
    return transformer_decode(
        self._decoder_function, decoder_input, encoder_output,
        encoder_decoder_attention_bias, decoder_self_attention_bias,
        hparams, attention_weights=self.attention_weights, cache=cache,
        decode_loop_step=decode_loop_step, nonpadding=nonpadding, losses=losses,
        **kwargs)

위 코드는 Transformer 클래스를 정의하는 부분인데요, 일단 일부만 가져왔습니다. 

 

이 코드를 보시면 Transformer 클래스는 t2t_model.T2TModel를 상속받습니다.

 

초기화를 해주고 encode와 decode 메서드를 정의하는데 

이 메서드는 앞서 정의한 transformer_encode()와 transformer_decode() 를 사용합니다. 

  def body(self, features):
    """Transformer main model_fn.

    Args:
      features: Map of features to the model. Should contain the following:
          "inputs": Transformer inputs. [batch_size, input_length, 1,
            hidden_dim].
          "targets": Target decoder outputs. [batch_size, decoder_length, 1,
            hidden_dim]
          "target_space_id": A scalar int from data_generators.problem.SpaceID.

    Returns:
      Final decoder representation. [batch_size, decoder_length, hidden_dim]
    """
    hparams = self._hparams

    losses = []

    if self.has_input:
      inputs = self._prepare_inputs_for_body(features)
      target_space = features["target_space_id"]
      encoder_output, encoder_decoder_attention_bias = self.encode(
          inputs, target_space, hparams, features=features, losses=losses)
    else:
      encoder_output, encoder_decoder_attention_bias = (None, None)

    targets = features["targets"]
    targets_shape = common_layers.shape_list(targets)
    targets = common_layers.flatten4d3d(targets)
    decoder_input, decoder_self_attention_bias = self._prepare_decoder_fn(
        targets, hparams, features=features)

    # Not all subclasses of Transformer support keyword arguments related to
    # recurrent memory, so only pass these arguments if memory is enabled.
    decode_kwargs = {}
    if self.recurrent_memory_by_layer is not None:
      # TODO(kitaev): The chunk_number feature currently has the same shape as
      # "targets", but this is only for the purposes of sharing sharding code.
      # In fact every token within an example must have the same chunk number.
      chunk_number_each_token = tf.squeeze(features["chunk_number"], (-1, -2))
      chunk_number_each_example = chunk_number_each_token[:, 0]
      # Uncomment the code below to verify that tokens within a batch share the
      # same chunk number:
      # with tf.control_dependencies([
      #     tf.assert_equal(chunk_number_each_token,
      #                     chunk_number_each_example[:, None])
      # ]):
      #   chunk_number_each_example = tf.identity(chunk_number_each_example)
      decode_kwargs = dict(
          recurrent_memory_by_layer=self.recurrent_memory_by_layer,
          chunk_number=chunk_number_each_example,
          )
    decoder_output = self.decode(
        decoder_input,
        encoder_output,
        encoder_decoder_attention_bias,
        decoder_self_attention_bias,
        hparams,
        nonpadding=features_to_nonpadding(features, "targets"),
        losses=losses,
        **decode_kwargs
        )
    expected_attentions = features.get("expected_attentions")
    if expected_attentions is not None:
      attention_loss = common_attention.encoder_decoder_attention_loss(
          expected_attentions, self.attention_weights,
          hparams.expected_attention_loss_type,
          hparams.expected_attention_loss_multiplier)
      return decoder_output, {"attention_loss": attention_loss}

    ret = tf.reshape(decoder_output, targets_shape)
    if losses:
      return ret, {"extra_loss": tf.add_n(losses)}
    else:
      return ret

이 body 함수는 Transformer 모델의 중심 기능을 담고 있습니다. 

 먼저, 초기화를 해줍니다. 

파라미터인 features 는 Transformer 모델에 대한 메타 데이터를 담고 있는 dictionary입니다. 

이 딕셔너리는 target을 가지고 있습니다. 

target은 학습 시에는 지도를 위한 정답이 들어있고 예측 시에는 아무것도 들어있지 않습니다. 

 

 

  1. 인코딩 : 
  • self.has_input이 True일 경우, 모델에 입력이 주어진 것으로 간주하고 인코딩 작업을 수행합니다.
  • 입력 데이터는 _prepare_inputs_for_body 함수를 통해 처리됩니다.
  • encoder_output과 encoder_decoder_attention_bias는 앞서 정의한 encode() 메서드로 얻습니다. 
  • 만약 self.has_input이 False인 경우, 인코더의 출력은 None으로 설정됩니다.

2. 디코딩 준비 :

targets에 features["target"]을 할당해줍니다. 

그리고 인코딩 과정과 마찬가지로 

decoder_input과 decoder_self_attention_bias를 _prepare_decoder_fn 으로 얻습니다. 

 

3. 디코딩:

  • decode 함수를 통해 디코딩 작업이 수행됩니다. 여기에서 앞서 얻은 인코더 출력, 주의 바이어스 및 기타 인자들이 사용됩니다.

4. 주의 손실(Attention Loss) 처리:

  • expected_attentions이 제공되는 경우, 인코더-디코더 간의 주의 손실을 계산하여 결과에 추가합니다.
  • expected_attentions이 제공된 경우, 모델의 실제 attention 가중치(self.attention_weights)와 비교하여 손실(attention_loss)이 계산됩니다. 이 손실은 모델의 학습을 안내하는 데 사용되며, 주어진 expected_attentions에 따라 attention 가중치를 조정하는 데 도움을 줍니다.
  • 자세한 사항은 common_attention.encoder_decoder_attention_loss를 확인해봐야 할 것 같습니다. 

이 뒤쪽으로도 Transformer 클래스의 메서드가 많이 정의됩니다. 

그런데, 전체 코드를 다 올리기엔 너무 긴 것 같아서 짧게 짧게 설명만 하도록 하겠습니다. 

 

  •  _prepare_inputs_for_body(self, features): features["inputs"] 을 리턴하여 인풋을 준비. 
  • _greedy_infer(self, features, decode_length, use_tpu=False): 디코드 하는 과정에서 가장 확률이 높은 1개만 출력
  •   def _beam_decode(self, features, decode_length, beam_size, top_beams, alpha, use_tpu=False):
    • beam search 방식으로 상위 beam_size 개수만큼 candidate 유지
  •   def _prepare_inputs_for_decode(self, features): 디코딩을 하기 전에 입력의 차원을 조절하는 역할
  • get_decode_end_id(self): 디코더가 처음으로 받는 초기 입력을 반환
  • def fast_decode(encoder_output,
                    encoder_decoder_attention_bias,
                    symbols_to_logits_fn,
                    hparams,
                    decode_length,
                    vocab_size,
                    init_cache_fn=_init_transformer_cache,
                    beam_size=1,
                    top_beams=1,
                    alpha=1.0,
                    sos_id=0,
                    eos_id=beam_search.EOS_ID,
                    batch_size=None,
                    force_decode_length=False,
                    scope_prefix="body/",
                    sampling_temperature=0.0,
                    top_k=-1,
                    cache=None): 이 함수는 인자가 굉장히 많습니다. fast_decode 함수는 디코딩을 본격적으로 하는 함수 입니다. 이 함수                      의 beam_size에 따라 _beam_decode()가 실행될지 greedy_infer()가 수행될지 결정됩니다. 이 함수 내부에서 while loop를 돌고 최종적으로 decoded_ids 라는 시퀀스 ID들과 scores에 각 시퀀스 ID에 대한 확률을 리턴합니다.

이것으로 transformer.py 중 Transformer 클래스의 정의 파트를 완료했습니다. 

 

이번 포스팅에서는 여기까지 살펴보도록 하고 

다음 포스팅에서 그 이후부터 함께 보도록 하겠습니다. 

 

감사합니다. 

 

블레이즈 테크노트.

저작자표시 비영리 변경금지 (새창열림)

'머신러닝(Machine Learning)' 카테고리의 다른 글

NLP 트랜스포머 코드 스터디 리뷰 (3) transformer.py  (0) 2023.08.21
NLP 트랜스포머 코드 스터디 리뷰 (2) transformer.py  (0) 2023.08.19
NLP 트랜스포머 모델 데이터셋 wmt14 다운로드하기  (0) 2023.08.05
NLP 트랜스포머 네 번째, 멀티 헤드 어텐션 Multi-Head Attention 알아보기  (0) 2023.08.01
NLP 트랜스포머 세 번째, 셀프 어텐션 알아보기  (0) 2023.07.30

관련글

티스토리툴바