TF2.0_LSTM_Seq2seq_BeamSearchDecoder_實戰教程

1. 序

此篇教程只有代碼實現，沒有理論部分。

適合有一定的理論基礎，對TF2.x有一些瞭解的人。如果不瞭解沒關係，傳送門：tensorflow2.0入門與實戰 2019年最通俗易懂的課程

如果看不懂可以在找找其他的相關視頻。

2. 代碼編寫

TF2.x具備熱執行(eager_execution)和自動建圖(auto-graph)兩個新的特性，與1.x的版本有很大的不同，提高了易用性。eager_execution允許我們像使用python一樣，寫一句執行一句。TF2.0給了非常不錯的模塊化，便於快速搭積木構建自己的模型，然後基於Auto-graph構建圖，即可運行，用戶更加傻瓜式操作。

Seq2seq主要包括encoder和decoder兩個部分。且一般在predict階段，decoder部分會重寫一個與train階段不同的decoder，但是其predict decoder的參數與train decoder的是一致的。所以，本片會寫兩個Seq2seq，一個是Seq2SeqTrainer，一個是Seq2SeqPredictor。

編寫的注意事項和各個部分作用的淺顯解釋，都在代碼的註釋中。尤其是Warning部分請一定看。

import tensorflow as tf
import tensorflow_addons as tfa


class Seq2SeqTrainer(tf.keras.Model):
    def __init__(self, batch_size, input_vocab_size, embedding_dims, Tx,
                 output_vocab_size, Ty, rnn_units, dense_layers
                ):
        """
        這裏個部分定義了模型的每層。也就是模型會用到哪些積木。
        
        Tx: 輸入到encoder的部分的句長
        Ty: decoder輸出的句長，也是Teacher forcus sequences長度
        """
        super(Seq2SeqTrainer, self).__init__()
        #########################################
        ##### encoder part
        #########################################
        ## embeddding layer: input word indices[batch_size, Tx] -- > encoder_word_embedding_mtx[batch_size, Tx, embedding_dims]
        self.encoder_embedding_fn = tf.keras.layers.Embedding(input_vocab_size, embedding_dims, Tx)
        ## LSTM layer: encoder_word_embedding_mtx --> rnn_output, rnn_a_tx, rnn_c_tx
        ## rnn_output: 每個時刻的LSTM單元的輸出的序列 [batch_size, Tx, rnn_units]
        ## rnn_a_tx: LSTM單元的cell state之一，[batch_size, rnn_units]
        ## rnn_c_tx: LSTM單元的cell state之一, [batch_size, rnn_units]
        self.encoder_rnn = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True)
        
        #########################################
        ##### decoder part
        #########################################
        ## embedding layer: decoder Teacher forcus sequences, word indices [batch_size, Ty] --> word_embedding_mtx [batch_size, Ty, embedding_dims]
        self.decoder_embedding_fn = tf.keras.layers.Embedding(output_vocab_size, embedding_dims, Ty)
        ## train sampler
        self.decoder_sampler = tfa.seq2seq.TrainingSampler()
        ## rnn_cells
        self.decoder_rnn_cells = tf.keras.layers.LSTMCell(rnn_units)
        self.decoder_dense_layer = tf.keras.layers.Dense(output_vocab_size)
        self.decoder_attention_mechanism = tfa.seq2seq.LuongAttention(dense_layers, None, batch_size*[Tx])
        self.decoder_rnn = tfa.seq2seq.AttentionWrapper(self.decoder_rnn_cells, self.decoder_attention_mechanism, attention_layer_size=dense_layers)
        self.decoder = tfa.seq2seq.BasicDecoder(self.decoder_rnn, sampler=self.decoder_sampler, output_layer=self.decoder_dense_layer)
        self.decoder_output_seq_length = batch_size*[Ty]
        
        ## other attri
        self.batch_size = batch_size
    
    def _decoder_initial_state(self, state):
        initial_state = self.decoder_rnn.get_initial_state(size=self.batch_size, dtype=tf.float32)
        return initial_state.clone(cell_state=state)

    def call(self, inputs, training=False):
        """
        Warning: 函數定義了模型的每層之間的銜接。也就是，如何組裝積木。告訴了程序如何組裝積木，它就會auto-graph生成正式的模型結構。
        
        seq2seq_trainer = Seq2SeqTrainer(...)
        seq2seq_trainer.compile(...)
        seq2seq_trainer.fit(...) # 在fit的時候，會自動調用這個函數，inputs是自己傳入的，training是程序自動給的training=True.
        
        seq2seq.predict(...) # 在預測時，也會自動調用這個函數，inputs是自己傳入的，training也是程序自動給的training=False.
        """
        encoder_inputs, deocder_inputs = inputs # 傳入的參數會自動轉化爲Tensor
        
        # encoding
        encoder_emb_inp = self.encoder_embedding_fn(encoder_inputs)
        a, a_tx, c_tx = self.encoder_rnn(encoder_emb_inp)
        
        # decoding
        decoder_emb_inp = self.decoder_embedding_fn(decoder_inputs)
        self.decoder_attention_mechanism.setup(a)
        outputs, _, __ = self.decoder(decoder_emb_inp, initial_state=self._decoder_initial_state([a_tx, c_tx]), sequence_length=self.decoder_output_seq_length)
        
        return outputs.rnn_output ## logits
        

## google 官方NMT模型中的loss
@tf.function
def loss_function(y, y_pred):
    sparsecategoricalcrossentropy = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True,
                                                                                  reduction='none')
    loss = sparsecategoricalcrossentropy(y_true=y, y_pred=y_pred)
    mask = tf.logical_not(tf.math.equal(y,0))   #output 0 for y=0 else output 1
    mask = tf.cast(mask, dtype=loss.dtype)
    loss = mask * loss
    loss = tf.reduce_mean(loss)
    return loss


seq2seq_trainer = Seq2SeqTrainer(...) # 填寫自己的參數
seq2seq_trainer.compile(optimizer="adam", loss=loss_function)

model_dir = "保存模型的文件夾路徑"
model_name = "你的模型的名字"
import os
model_name_prefix = os.path.join(model_dir, "%s_epoch{epoch}" % model_name)
checkpnt_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=model_name_prefix,
    save_weights_only=True ### Warning：繼承於tf.keras.Model, 定義的模型的實例只能保存模型參數，不可以保存整個模型，
)
seq2seq_trainer.fit(dataset, epochs=5, callbacks=[checkpnt_callback])

示例輸出：

Train for 1474 steps
Epoch 1/5
1474/1474 [==============================] - 896s 608ms/step - loss: 1.8931
Epoch 2/5
1474/1474 [==============================] - 891s 605ms/step - loss: 1.1537
Epoch 3/5
1474/1474 [==============================] - 891s 604ms/step - loss: 0.8433
Epoch 4/5
1474/1474 [==============================] - 894s 607ms/step - loss: 0.6514
Epoch 5/5
1474/1474 [==============================] - 894s 606ms/step - loss: 0.5161

3. 部分Bug

待更新