Tensorflow_addons seq2seq example using Attention and Beam Search

이글은 medium.com의 내용임. (링크)

Tensorflow_addons seq2seq example using Attention and Beam Search

    텐서플로는 neural machine translation을 나타내는 nmt 저장소(repository)를 호스팅하고 어텐션 기반 인코더-디코더 seq2seq 모델을 사용하는 방법을 제공한다. 어텐션 모델을 구현하는 API는 tf.contrib 모듈로부터 제공된다. TF2.0에서 tf.contrib 모듈로부터의 리소스는 tensorflow_addon(tfa) 또는 tensorflow_io같은 구분된 프로젝트로 이동되었다. nmt에 관련된 contrib 리소스 또한 재설계되었다. 이것은 tensorflow core의 부분이 아니기 때문에 추가적으로 설치해야 한다.

pip install tensorflow-addons

https://github.com/tensorflow/nmt ( TF 1.14)

TensorFlow Addons는 잘 정립된 API 패턴을 따르지만 core TensorFlow에서 가능하지 않은 새로운 기능을 구현하는 기여 저장소이다. 텐서플로는 기본적으로 많은 연산(operator), 레이어, 지표(metrics) 그리고 옵티마이저를 제공한다. 그러나, ML 같은 빨르게 움직이는 분야에서 텐서플로 코어에 통합되지 못하는 많은 흥미로운 개발이 있다.(넓은 응용 가능성이 아직 명화하지 않거나 커뮤니티의 더 작은 하위집단에서 주로 사용되기 때문이다.)

    nmt에 대한 텐서플로 공식 튜토리얼은 애드온 API에 기초하지 않는다. 이것은 설계부터 스스로 어텐션 로직을 어떻게 구축하는지를 보여준다. 예를 들면 다음 링크는 Bahdanau 어텐션에 기초한 NMT에 대한 튜토리얼 이다.

Neural machine translation with attention

    이 글의 목적은 tfa.seq2seq API를 사용하여 구축된 모델에 대한 설계, 훈련 그리고 추론을 포함하는 단데단 튜토리얼을 공유하는 것이다. 어텐션(attention)과 빔써치(beam search) 디코딩을 기반으로 영어를 프랑스어로 번역하는 모델을 구축해보자.

ENCODRE

    인코딩 네트워크는 seq2seq API에 의존하지 않는다. 인코딩 네트워크는 동일한 tf.keras API로 구현된다. LSTM 레이어는 입력 임베딩 배치로 인코딩을 출력한다(seq2seq API에서 메모리라고도 함.) 또한 최종단계의 hidden activation과 메모리를 출력한다.

Encoder Network

#ENCODER Network

class EncoderNetwork(tf.keras.Model):
    def __init__(self, input_vocab_size, embedding_dims, rnn_units):
        super().__init__()
        self.encoder_embedding = tf.keras.layers.Embedding(input_dim=input_vocab_size, output_dim=embedding_dims)
        self.encoder_rnnlayer = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True)

DECODER with ATTENTION MECHANISM

    디코더 네트워크는 디코더와 어텐션 메커니즘 모두를 포함한다. 어텐션 메커니즘 객체(AttentionMechanism object)는 Bahdanau/Luong 같은 어텐션 메커니즘의 상세 구현을 포함한다. 어텐션 메커니즘은 인코더 메모리(코드에서 'a'로 표시)와 이전 디코더 hidden state([a_tx]로 표시) 모두에 접근한다.

    어텐션랩퍼 객체(AttentionWrapper object)는 디코더 LSTM cell과 어텐션 메커니즘 객체의 참조를 사용하여 생성된다. 어텐션매커니즘 객체는 사용되기 전에 인코더 메모리로 초기화 되어야 한다. 그러나 여기서는 훈련동안 배치에서 동적으로 메모리를 전달해야 하기 때문에 디코더 네트워크 클래스 정의에서 메모리 없이 이 객체를 생성한다. 따라서 디코더 메모리는 'none'으로 전달된다.

    LSTM cell은 어텐션 메커니즘이 각 단계에서 다른 어덴션 가중치를 생성해야 하기 때문에 디코더에서 사용된다. LSTM cell은 또한 이전 hidden과 메모리 상태 (a_tx, c_tx)로 초기화 되어야 한다.

    기본 디코더 객체(BasicDecoder object)는 어텐션랩퍼, 밀집(dense) 레이어 그리고 훈련 샘플러(training sampler)의 참조를 갖는다. 기본 디코더는 LSTM cell을 초기화하기 위해 기본 디코더의 initialize() 메소드를 한번 호출하고 각 단계에서 되풀이하여 step() 함수를 호출한다.

#Decoder Network with Attention

class DecoderNetwork(tf.keras.Model):
    def __init__(self, output_vocab_size, embedding_dims, rnn_units):
        super().__init__()
        self.decoder_embedding = tf.keras.layers.Embedding(input_dim=output_vocab_size, output_dim=embedding_dims)
        self.dense_layer = tf.keras.layers.Dense(output_vocab_size)
        self.decoder_rnncell = tf.keras.layers.LSTMCell(rnn_units)
        #Sampler
        self.sampler = tfa.seq2seq.sampler.TrainingSampler()
        # memory=None인 어텐션 메커니즘 생성
        self.attention_mechanism = self.build_attention_mechanism(dense_uinits, None, BATCH_SIZE * [Tx])
        self.rnn_cell = self.build_rnn_cell(BATCH_SIZE)
        self.decoder = tfa.seq2seq.BasicDecoder(self.rnn_cell, sampler=self.sampler, output_layer=self.dense_layer)

    def build_attention_mechanism(self, units, memory, memory_sequence_length):
        return tfa.seq2seq.LuongAttention(units, memory=memory, memory_sequence_length=memory_sequence_length)

        #return tfa.seq2seq.BahdanauAttention(units, memory=memory, memory_sequence_length=memory_sequence_length)

    # wrap decodernn cell
    def build_rnn_cell(self, batch_size):
        rnn_cell = tfa.seq2seq.AttentionWrapper(self.decoder_rnncell, self.attention_mechanism, attention_layer_size=dense_units)
        return rnn_cell

    def build_decoder_initial_state(self, batch_size, encoder_state, Dtype):
        decoder_initial_state = self.rnn_cell.get_initial_state(batch_size=batch_size, dtype=Dtype)
        decoder_initial_state = decoder_initial_state.clone(cell_state=encoder_state)
        return decoder_initial_state

encoderNetwork = EncoderNetwork(input_vocab_size, embedding_dims, rnn_units)
decoderNetwork = DecoderNetwork(output_vocab_size, embedding_dims, rnn_units)

TRAINING

    훈련의 한 단계는 인코더 네트워크로부터의 출력을 가져와 디코더 네트워크에 전달한다.

#One step of training on a batch using Teacher Forcing techique

def train_step(input_batch, output_batch, encoder_initial_cell_state):
    #loss = 0으로 초기화
    loss = 0
    with tf.GradientTape() as tape:
        encoder_emb_inp = encoderNetwork.encoder_embedding(input_batch)
        a, a_tx, c_tx = encoderNetwork.encoder_rnnlayer(encoder_emb_inp, initial_state=encoder_initial_cell_state)

        decoder_input = output_batch[:, :-1] # end 무시
        # decoder_input의 timestep + 1 버전과 로짓 비교
        decoder_output = output_batch[:, 1:] # start 무시

        # Decoder Embedding
        decoder_emb_inp = docoderNetwork.decoder_embedding(decoder_input)
        # 디코더 출력과 AttentionWrapperState에 대한 zero state로 부터 디코더 메모리 셋업
        decoderNetwork.attention_mechanism.setup_memory(a)
        # 디코더 네트워크에 입력으로써 전달된 인코더의 [last step activation, last memory_stat]
        decoder_initial_state = decoderNetwork.build_initial_state(BATCH_SIZE, encoder_state=[a_tx, c_tx], Dtype=tf.float32)

        # BasicDecoderOutput
        outputs, _, _ = decoderNetwork.decoder(decoder_emp_inp, initial_state=decoder_initial_state, sequence_length=BATCH_SIZE*[Ty-1])

        logits = ouputs.rnn_output

        # 손실 계산
        loss = loss_function(logits, decoder_output)

    # 모든 레이어의 변수/가중치 목록 반환
    variables = encoderNetwork.trainable_variables + decoderNetwork.trainable_variables
    # 손실 wrt 변수 미분
    gradients = tape.gradient(loss, variables)

    # grads_and_vars - list of (gradient, variable) pairs
    grads_and_vars = zip(gradient, variables)
    optimizer.apply_gradients(grads_and_vars)

    return loss

    decoder_initial-state가 [s_prev]인 Vanilla 인코더-디코더(Attention 없이)와 달리, 여기에서는 아래 텐서를 갖는 AttentionWrapperState 객체이다.

AttentionWrapperState Object

inference_batch_size = 2, rnn_units = 1024, embedding_di = 256

INFERENCE with greedy sampling

    모델로부터 추론(inference)을 생성하는 것이 모텔을 훈련하는 것과 차이가 있기 때문에 또다른 BasicDecoder를 만든다. 샘플링동안 사용가능한 입력은 오직 입력 인코딩과 디코딩 네트워크를 동작시키기 위한 <start> 레이블 뿐이다. 디코더는 <start>레이블을 따르는 다음 레이블 셋 예측을 시작해야 한다. 디코더는 greedy sampling 또는 beam search decoding 방법을 구현할 수 있다. 훈련 단계에서 전체 디코더 입력은 모든 단계에서 가능하다. 따라서 training sampler가 사용된다. 추론 단계에서는 디코더에 대한 다음 단계가 이전 단계로부터의 예측을 사용하여 구성기 때문에 greedy sampler가 사용된다. 또한 디코더 객체는 수동으로 단계가 진행되어야 한다.

# Inference with Greedy Sampling

encoder_initial_cell_state = [tf.zeros((inference_batch_size, rnn_units)), tf.zeros((inference_batch_size, rnn_units))]
encoder_emb_inp = encoderNetwork.encoder_embedding(inp)
a, a_tx, c_tx = encoderNetwork.encoder_rnnlayer(encoder_emb_inp, initial_state=encoder_initial_cell_state)

start_tokens = tf.fill([inference_batch_size], Y_tokenizer.word_index['<start>'])
end_token = Y_tokenizer.word_index['<end>']
greedy_sampler = tfa.seq2seq.GreedyEmbeddingSampler()

decoder_input = tf.expand_dims([Y_tokenizer.word_index['<start>]]) * inference_batch_size, 1)
decoder_emb_inp = decoderNetwork.decoder_embedding(decoder_input)

decoder_instance = tfa.seq2se1.BasicDecoder(cell = decoderNetwork.rnn_cell, sampler= greedy_sampler, output_layer=decoderNetwork.dense_layer)
decoderNetwork.attention_mechanism.setup_memory(a)
# LSTM에 대해 입력으로 [Lsat step activations, encoder_memory_state]를 디코더에 전달
print("decoder_initial_state = [a_tx, c_tx] :", np.array([a_tx, c_tx]).shape)
decoder_initial_state = decoderNetwork.build_decoder_initial_state(inference_batch_size, encoder_state=[a_tx, c_tx], Dtype=tf.float32)

maximum_iterations = tf.round(tf.reduce_mean(Tx) * 2)

# 추론 디코더 초기화
(first_finished, first_inputs, first_state) = decoder_instance.initialize(decoder_embedding_matrix, start_tokens=start_tokens, end_token=end_tokens, initial_state=decoder_initial_state)

inputs = first_inputs
state = first_state
predictions = np.empty((inference_batch_size, 0), dtype=np.int32)
for j in range(maximum_iterations):
    outputs, next_state, next_inputs, finished = decoder_instance.step(j, inputs, state)
    inputs = next_inputs
    state = next_state
    outputs = np.expand_dims(outputs.sample_id, axis = -1)
    predictions = np.append(predictions, outputs, axis = -1)

    첫번째 시퀀스에서의 <end>에서 번역하지 않음.

INFERENCE with BEAM SEARCH

    빔써치 디코더(BeamSearchDecoder)는 BasicDecoder에 대한 경우인 단일 sample_id 대신 predicted_ids를 출력한다.(tfa.seq2seq.BeamSearchDecoderOutput) 점수는 현재 단계에 완전히 합산되어지고 더 높은 점수를 생성하는 시퀀스가 빔 서치 출력을 차지한다. 빔써치가 동작하기 위해 존재하는 추론 블록(inference block)은 매우 작은 변화만이 필요하다. 인코더 출력은 빔의 폭(beam width)에 맞춰지고 디코더 네늩워크로 전달된다. 유사하게 어텐션랩퍼 객체 또한 입력 배치크기와 빔의 폭을 곱하여 초기화된다.

# BeamSearchDecoder

# from official documentation
# NOTE if you ar using the BeamSearchDecoder with a cell wrapped in AttentionWrapper, then you must ensure that:

# The encoder output has been tiled to beam_width via tfa.seq2se1.tile_batch (NOT tof.tile)
# The batch_size argument passed to the get_initial_state method of this wrapper is equal to true_batch_size * beam_width
# the inital state created with get_initial_state above contains a cell_state value containing properly tiled final state form the encoder.
encoder_memory = tfa.seq2seq.tile_batch(a, beam_width)
decoderNetwork.attention_mechanism.setup_memory(encoder_memory)
print("beam_width * [batch_size, Tx, rnn_units] : 3 * [2, Tx, rnn_units] : ", encoder_memory.shape)
# set decoder_initial_state which is an AttentionWrapperState considering beam_width
decoder_inital_state = decoderNetwork.rnn_cell.get_initial_state(batch_size=inference_batch_size * beam_width, dtype=Dtype)
encoder_state = tfa.seq2seq.tile_batch(s_prev, muliplier=beam_width)
decoder_initial_state = decoder_initial_state.clone(cell_state=encoder_state)

decoder_instance = tfa.seq2seq.BeamSearchDecoder(decoderNetwork.rnn_cell, beam_width=beam_width, output_layer=decoderNetwork.dense_layer)

    가장 높은 점수 (-0.288, -7.808)인 빔써치 출력은 영어 입력에 일치하기 위한 올바른 프랑스어 번역이다.