TF2.0_LSTM_Seq2seq_BeamSearchDecoder_实战教程

1. 序

此篇教程只有代码实现，没有理论部分。

适合有一定的理论基础，对TF2.x有一些了解的人。如果不了解没关系，传送门：tensorflow2.0入门与实战 2019年最通俗易懂的课程。如果看不懂可以找找其他的相关视频，毕竟广听不错。

修BUG的方法：先Google，再官方文档，不行看源码。官方文档的很多函数都有其对应的github源码链接，非常方便。本人对于TF1.x仅仅用过一丢丢，但是看它2.x的代码感觉还好，不难，易看懂。

另一外的秘密是：虽然TF2.x与TF1.x差别大，但是貌不似神似，即，代码不互通，逻辑互通的。
悲惨的一点：貌似TF2.0用户不多，例子也不多。

才疏学浅，多多斧正。

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部分请一定看。

2.1 training部分

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, 定义的模型的实例只能保存模型参数，不可以保存整个模型，
)
## dataset请自己按照其他的教程实现吧，比如这个(小美女的教程)[https://zhuanlan.zhihu.com/p/59506402]
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

2.2 predicting部分

class Seq2SeqPredictForSave(tf.keras.Model):
    def __init__(self, Ty, beam_width, start_token, end_token, seq2seq_train):
        super(Seq2SeqPredictForSave, self).__init__()
        ############################
        ###   encoder部分的积木
        ###########################
        self.encoder_embedding = seq2seq_train.encoder_embedding
        self.encoder_rnn = seq2seq_train.encoder_rnn
        
        ############################
        ###   decoder部分的积木
        ###########################
        self.decoder_embedding_mtx = seq2seq_train.decoder_embedding.variables[0]
        self.decoder_attention_mechanism = seq2seq_train.decoder_attention_mechanism
        self.decoder_rnn_cell = seq2seq_train.decoder_rnn_cell
        ## beam_search_decoder
        self.beam_search_decoder = tfa.seq2seq.BeamSearchDecoder(
            cell=self.decoder_rnn_cell,
            beam_width=beam_width,
            embedding_fn=seq2seq_train.decoder_embedding, ## 一定，一定，以这个为准。对于输入的start_token,end_token和预测过程中产生的word token做lookup转换。
            output_layer=seq2seq_train.decoder_dense_layer
        )
        
        ############################
        ### 其他属性
        ###########################
        self.beam_width = beam_width
        ## 在google的NMT教程中给的值是大约 Ty * 2
        self.maximum_iterations = Ty * 2
        self.start_token = start_token
        self.end_token = end_token
        
    def _decoder_initial_state(self, size, state):
        ## Warning: 此处的dtype需要与你在训练模型中的dtype一致。
        _initial_state = self.decoder_rnn_cell.get_initial_state(batch_size=size*self.beam_width, dtype=tf.float32)
        decoder_initial_state = _initial_state.clone(cell_state=state)
        return decoder_initial_state
    
    def _generate_lines(self, start_tokens, tiled_state):
        beam_width = self.beam_width
        (first_finished, first_inputs, first_state) = self.beam_search_decoder.initialize(embedding=None, ## Warning: 这个地方不要传值过来。这个参数与构造函数中的"embedding_fn"是同一个东西，且二者不能同时都传值，否则报错。如果，你使用这个参数，而非"embedding_fn", 你的代码只能predict一次，第二次就会报错，这是他们逻辑上的bug。
                                                                                         start_tokens=start_tokens,
                                                                                         end_token=self.end_token,
                                                                                         initial_state=self._decoder_initial_state(1, state=tiled_state),
                                                                                        )
        inputs = first_inputs 
        state = first_state
        # 最终会生成beam_width条最有路径，它们用两个矩阵盛放。
        # pred_mtx t时刻，预测的beam_width个最好的句子的当前最佳单词
        # path_mtx 句子的t-1时刻的单词对应的单词位置
        # 表述不清，建议输出这两个矩阵看看。
        pred_mtx, path_mtx = [], []

        for j in range(self.maximum_iterations):
            outputs, next_state, next_inputs, finished = self.beam_search_decoder.step(j, inputs, state)
            inputs = next_inputs
            state = next_state

            pred_mtx.append(outputs.predicted_ids)
            path_mtx.append(outputs.parent_ids)
        
        return pred_mtx, path_mtx

    def call(self, inputs, training=False):
        # 输入
        # Warning：一次只能预测一个样本
        encoder_input = inputs
        #########################
        ## 组装encoder部分的积木
        #########################
        encoder_x = self.encoder_embedding(encoder_input)
        a, a_tx, c_tx = self.encoder_rnn(encoder_x)
        
        beam_width = self.beam_width
        ###########################
        ## tilding部分  铺砖
        ###########################
        tiled_a = tfa.seq2seq.tile_batch(a, multiplier=self.beam_width)
        tiled_a_tx = tfa.seq2seq.tile_batch(a_tx, multiplier=self.beam_width)
        tiled_c_tx = tfa.seq2seq.tile_batch(c_tx, multiplier=self.beam_width) 
        # 一次只能预测一个句子，batch_size = 1
        start_tokens = tf.fill([1], self.start_token)
        
        ##########################
        ## 组装decoder部分的积木
        ##########################
        self.decoder_attention_mechanism.setup_memory(tiled_a)
        pred_mtx, path_mtx = self._generate_lines(start_tokens, [tiled_a_tx, tiled_c_tx])
        # Warning:
        # pred_mtx, path_mtx 会被按照行拼接为mtx矩阵，再输出
        # pred_mtx [self.maximum_iterations, beam_width]
        # path_mtx [self.maximum_iterations, beam_width]
        # mtx [self.maximum_iterations * 2, beam_width]
        # pred_mtx 对应mtx的前self.maximum_iterations行
        return pred_mtx, path_mtx

# 这段程序主要是为了将预测出来的数据是一个矩阵，将矩阵中的句子组装出来。
# 没有在SeqSeqPredictor中组装，是因为我的对TF不熟悉。见谅！
# 如果您对TF熟悉，给出您的方案。我会添加，并著名您是作者。
import numpy as np
def predict(test_x, test_y):
    beam_width, maximum_iterations = seq2seq_predictor.beam_width, seq2seq_predictor.maximum_iterations
    
    mtx = seq2seq_predictor.predict(test_x)
    pred_mtx = [arr[0] for arr in mtx[:maximum_iterations]]
    path_mtx = [arr[0] for arr in mtx[maximum_iterations:]]
    
    pred_lines = np.zeros((beam_width, maximum_iterations), dtype=np.int32)
    pred_lines[:, -1] = pred_mtx[-1]

    path_mtx.pop(0)

    pre_pred_lines = np.array([[pred[p] for p in path] for pred, path in zip(pred_mtx, path_mtx)])

    pred_lines[:, 0:-1] = pre_pred_lines.T
	
	# idx是index的缩写，opword是output word的缩写
	# idx2opword, 是一个np.array([output_word1, output_word2, output_word3, ..., output_wordn])
	# 生成的句子中的词的index，Seq2Seq模型的输出句子的单词。
	# 不懂的参照这个(小美女的教程)[https://zhuanlan.zhihu.com/p/59506402]
    predict_sents = [idx2opword[l] for l in pred_lines]

    print("# input:")
    # 与idx2opword类似，不过ipword是input word，输入seq2seq模型的句子的单词。
    print(" ".join(idx2ipword[test_x[0]]))

    print("\n# Standard:")
    print(" ".join(idx2opword[test_y[0]]))

    for i, l in enumerate(predict_sents):
        print(f"##{i}--------------------------")
        line = []
        for w in l:
            if w == "":
                break
            line.append(w)
        print(" ".join(line))

save和load

待续，好长呀，一个小时没了，明天再更。

3. 部分Bug

待更新