머신러닝(Machine Learning)
NLP 트랜스포머 코드 스터디 리뷰 (1) transformer.py
Blaze_블즈
2023. 8. 18. 02:06
안녕하세요
블레이즈 테크노트
블레이즈 입니다.
트랜스포머 논문에 대해서는,
제가 지난 여러 포스팅에서 설명했습니다.
다시 짚고 넘어가자면 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라는 함수의 정의가 시작됩니다.
- 입력 매개변수:
- encoder_function: 실제 인코딩을 수행하는 함수입니다
- inputs: 인코딩이 필요한 초기 입력 텐서입니다.
- target_space: 대상 공간을 식별합니다
- hparams: 모델의 하이퍼파라미터입니다.
- 기타 매개변수들은 주의 깊게 처리하거나 저장할 필요가 있는 추가 정보를 위해 사용됩니다.
- 전처리:
- 입력은 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은 학습 시에는 지도를 위한 정답이 들어있고 예측 시에는 아무것도 들어있지 않습니다.
- 인코딩 :
- 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 클래스의 정의 파트를 완료했습니다.
이번 포스팅에서는 여기까지 살펴보도록 하고
다음 포스팅에서 그 이후부터 함께 보도록 하겠습니다.
감사합니다.
블레이즈 테크노트.