tensorflow2.0官网demo学习笔记 基于attention的seq2seq机器翻译

前言

打算详细深入了解bert和transform的原理和源代码，先从基于attention的seq2seq学起，代码来自tf2.0官网demo，个人修正了一部分bug，可以跑通，并做了更详尽的注释。

备注

1. 西班牙语-英语的翻译，自己换个训练文件也能跑通
2. 参考自官方demo，网址https://www.tensorflow.org/tutorials/text/nmt_with_attention
3. 原代码有部分bug，已修正，保证跑通
4. 注释更详尽，已尽力做到傻瓜式

代码

import os
# os.environ["CUDA_VISIBLE_DEVICES"]="-1"   
# 本地调试我用的核显本，无cuda，故屏蔽GPU
import tensorflow as tf

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from sklearn.model_selection import train_test_split

import unicodedata
import re
import numpy as np
import io
import time
if gpus:
   try:    
       for gpu in gpus:
           tf.config.experimental.set_memory_growth(gpu, True)
   except RuntimeError as e:
       print(e)
#设置增长式显存占用

备注

下载和准备数据集
我们将使用 http://www.manythings.org/anki/ 提供的一个语言数据集。这个数据集包含如下格式的语言翻译对：

May I borrow this book? ¿Puedo tomar prestado este libro?
这个数据集中有很多种语言可供选择。我们将使用英语 - 西班牙语数据集。为方便使用，我们在谷歌云上提供了此数据集的一份副本。但是你也可以自己下载副本。下载完数据集后，我们将采取下列步骤准备数据：

1、给每个句子添加一个 开始 和一个 结束 标记（token）。
2、删除特殊字符以清理句子。
3、创建一个单词索引和一个反向单词索引（即一个从单词映射至 id 的词典和一个从 id 映射至单词的词典）。
4、将每个句子填充（pad）到最大长度。

可能是因为语言数据集更新了的缘故，实际上语言翻译对的格式不是上述的两个句子，而是三个句子，第三个句子对翻译任务来说没什么作用，下面的代码注释里会写明。

# 下载文件
# path_to_zip = tf.keras.utils.get_file(
#     'spa-eng.zip', origin='http://storage.googleapis.com/download.tensorflow.org/data/spa-eng.zip',
#     extract=True)
# path_to_file = os.path.dirname(path_to_zip)+"/spa-eng/spa.txt"

path_to_file="./spa-eng/spa.txt"
# 因为我自己下好了文件，直接指定文件位置就行。直接用上述代码下载的话，可能需要

def unicode_to_ascii(s):
   """
   将 unicode 文件转换为 ascii

   unicodedata.normalize：把一串UNICODE字符串转换为普通格式的字符串，具体格式支持NFC、NFKC、NFD和NFKD格式。
   如果需要一种单一的单一的表示方式，可以使用一种规范化的Unicode文本形式来减少不想要的区别。
   Unicode标准定义了四种规范化形式： Normalization Form D (NFD)，Normalization Form KD (NFKD)，Normalization Form C (NFC)，和Normalization Form KC (NFKC)。
   大约来说，NFD和NFKD将可能的字符进行分解，而NFC和NFKC将可能的字符进行组合。

   unicodedata.category：把一个字符返回它在UNICODE里分类的类型。
   
   参考自https://www.cnblogs.com/wz123/archive/2004/01/13/12303932.html
   不同的编码方式覆盖的字符集不一样，尤其是这种涉及到不同国家语言的，转码之后，各字符的表示会统一。
   照着用就完事了，拿去做文本生成什么的任务，也用不着转码。
   :param s: 转码前的文本
   :return: 转码后的文本
   """
   return ''.join(c for c in unicodedata.normalize('NFD', s)
       if unicodedata.category(c) != 'Mn')

def preprocess_sentence(w):
   w = unicode_to_ascii(w.lower().strip())

   # 在单词与跟在其后的标点符号之间插入一个空格
   # 例如： "he is a boy." => "he is a boy ."
   # 参考：https://stackoverflow.com/questions/3645931/python-padding-punctuation-with-white-spaces-keeping-punctuation
   w = re.sub(r"([?.!,¿])", r" \1 ", w)
   w = re.sub(r'[" "]+', " ", w)

   # 除了 (a-z, A-Z, ".", "?", "!", ",")，将所有字符替换为空格
   w = re.sub(r"[^a-zA-Z?.!,¿]+", " ", w)

   w = w.rstrip().strip()

   # 给句子加上开始和结束标记
   # 以便模型知道何时开始和结束预测
   w = ' ' + w + ' '
   return w        

现在来看一看上面俩函数什么效果

en_sentence = u"May I borrow this book?"
sp_sentence = u"¿Puedo tomar prestado este libro?"
print(preprocess_sentence(en_sentence))
print(preprocess_sentence(sp_sentence).encode('utf-8'))

# 为方便阅读，直接把输出附上：
# <start> may i borrow this book ? <end>
# b' \xc2\xbf puedo tomar prestado este libro ? '

def create_dataset(path, num_examples):
   """
   原文的注释说的函数作用是：
   1. 去除重音符号
   2. 清理句子
   3. 返回这样格式的单词对：[ENGLISH, SPANISH]
   我对第三点存疑，因为很明显，下面的代码会返回三个元素的可迭代对象，即[ENGLISH, SPANISH, 第三个不知道啥(可以看函数内的注释)]
   :param path: 文件地址
   :param num_examples:读取多少组翻译对照，其实就是训练数据集大小
   :return: 返回的翻译对
   """
   lines = io.open(path, encoding='UTF-8').read().strip().split('\n')
   # lines的一行大概是 'Who?\t¿Quién?\tCC-BY 2.0 (France) Attribution: tatoeba.org #2083030 (CK) & #2122720 (Shishir)'
   # 下面的列表迭代式中，按\t分割，就能分成三段，形如：
   # ['Who?', '¿Quién?', 'CC-BY 2.0 (France) Attribution: tatoeba.org #2083030 (CK) & #2122720 (Shishir)']
   word_pairs = [[preprocess_sentence(w) for w in l.split('\t')]  for l in lines[:num_examples]]

   return zip(*word_pairs)

测试下这个函数：

en, sp, _ = create_dataset(path_to_file, None)
# 原文是en, sp = create_dataset(path_to_file, None), 因为函数返回三个值，我这里添加了个占位的_，不然报错
print(en[-1])
print(sp[-1])

# 输出：
# <start> if you want to sound like a native speaker , you must be willing to practice saying the same sentence over and over in the same way that banjo players practice the same phrase over and over until they can play it correctly and at the desired tempo . <end>
# <start> si quieres sonar como un hablante nativo , debes estar dispuesto a practicar diciendo la misma frase una y otra vez de la misma manera en que un musico de banjo practica el mismo fraseo una y otra vez hasta que lo puedan tocar correctamente y en el tiempo esperado . <end>

我觉得这个函数不需要注释：

def max_length(tensor):
   return max(len(t) for t in tensor)

def tokenize(lang):
   lang_tokenizer = tf.keras.preprocessing.text.Tokenizer(
     filters='')
   # Tokenizer是一个用于向量化文本的类，即单词在整个文本库中的下标构成的列表，从1开始。更多介绍在函数下方有注释。
   lang_tokenizer.fit_on_texts(lang)
   # fit_on_texts方法：传入要用以训练的文本

   tensor = lang_tokenizer.texts_to_sequences(lang)
   # texts_to_sequences方法：传入待转为序列的文本列表，返回文本序列列表

   tensor = tf.keras.preprocessing.sequence.pad_sequences(tensor,
                                                        padding='post')
   # 序列填充，超过最大长度(默认maxlen=None)的截断，不传参的话，默认照序列列表的最大长度，
   # 因为这里主要是短句翻译，所以默认参数就行，要是做文本分类的话，还是限制下序列长度比较好。
   # 小于最大长度的补足，不传参的话默认补0，padding表示补足方式，post是末尾补足(默认是pre，表示首位补0)
   # 更多参数自行查看api手册
   return tensor, lang_tokenizer 
   # tensor：张量，文本序列列表
   # lang_tokenizer：已经传入训练文本的实例化的Tokenizer类

关于tf.keras.preprocessing.text.Tokenizer类的举例：
texts = [“你好 我好 你好 你好 你好 我们 大家 都 好 吗 吗 吗 吗 吗”, “分词器 训练 文档 训练 文档 文档 你好 我好”]
lang_tokenizer = tf.keras.preprocessing.text.Tokenizer(filters=’’)
lang_tokenizer.fit_on_texts(texts)
一番操作，最后lang_tokenizer.texts_to_sequences(texts)会返回[[1, 4, 1, 1, 1, 6, 7, 8, 9, 2, 2, 2, 2, 2], [10, 5, 3, 5, 3, 3, 1, 4]]
两段文本，一共出现过10个词，最后用1表示你好，4是我好，很容易理解

def load_dataset(path, num_examples=None):
   # 创建清理过的输入输出对
   targ_lang, inp_lang, _ = create_dataset(path, num_examples)
   # 我加了个占位符，不然报错，参考测试create_dataset()函数时的注释
   input_tensor, inp_lang_tokenizer = tokenize(inp_lang)
   # 不清楚这行是在干嘛的话，看看上面tokenize()函数的注释
   target_tensor, targ_lang_tokenizer = tokenize(targ_lang)
   return input_tensor, target_tensor, inp_lang_tokenizer, targ_lang_tokenizer

限制数据集的大小以加快实验速度（可选）
在超过 10 万个句子的完整数据集上训练需要很长时间。为了更快地训练，我们可以将数据集的大小限制为 3 万个句子（当然，翻译质量也会随着数据的减少而降低）：

num_examples = 30000
# 设置数据集大小
input_tensor, target_tensor, inp_lang, targ_lang = load_dataset(path_to_file, num_examples)
# 这里输入语言和输出语言是分开编码的（逐级查找函数调用就可以理解，load_dataset里调用了两次tokenizer函数，两个Tokenizer类传入了不同的训练文本）

max_length_targ, max_length_inp = max_length(target_tensor), max_length(input_tensor)
# 计算目标张量的最大长度 （上面写的函数max_length）
# 本例中，max_length_targ, max_length_inp分别是11和16

input_tensor_train, input_tensor_val, target_tensor_train, target_tensor_val = train_test_split(input_tensor, target_tensor, test_size=0.2)
# 采用 80 - 20 的比例切分训练集和验证集。train_test_split是来自sklearn的数据集划分工具，大家都爱用。

print(len(input_tensor_train), len(target_tensor_train), len(input_tensor_val), len(target_tensor_val))
# 看一眼数据集长度

# 输出：
# 24000 24000 6000 6000

def convert(lang, tensor):
   """
   这就是个查看序列化文本有没有问题的函数，可以看下面的调用示例，输入一条序列，然后用lang.index_word方法看看对应的单词是什么。
   这里的lang，由函数load_dataset产生，应该叫lang_tokenizer，是个实例化的Tokenizer类。虽然变量名可以随意起，但是demo里这么写容易让人搞混。
   :param lang: 已经传入训练文本的Tokenizer类
   :param tensor: 一行序列化的文本，例如：array([  1,  20,  27,  43, 187,   3,   2,   0,   0,   0,   0,   0,   0, 0,   0,   0])
   :return: 
   """
   for t in tensor:
       if t!=0:
           print ("%d ----> %s" % (t, lang.index_word[t]))

# 上面的convert函数调用，可以看到序列化文本还原后是什么句子。
print ("Input Language; index to word mapping")
convert(inp_lang, input_tensor_train[0])
print ()
print ("Target Language; index to word mapping")
convert(targ_lang, target_tensor_train[0])

# 输出：
# Input Language; index to word mapping
# 1 ----> <start>
# 12 ----> me
# 808 ----> gustaba
# 86 ----> mucho
# 14 ----> la
# 911 ----> tarta
# 3 ----> .
# 2 ----> <end>

# Target Language; index to word mapping
# 1 ----> <start>
# 4 ----> i
# 26 ----> was
# 1078 ----> fond
# 61 ----> of
# 352 ----> cake
# 3 ----> .
# 2 ----> <end>

创建一个 tf.data 数据集：

BUFFER_SIZE = len(input_tensor_train)  # 多少条输入数据
BATCH_SIZE = 64
steps_per_epoch = len(input_tensor_train)//BATCH_SIZE  
# 字面上看就是整个数据集能取几次batch，所以他这个变量名翻译过来是每个epoch能跑几个step，一个step就是一个batch，也挺好理解
embedding_dim = 256  # word2vector维度
units = 1024  # 其实就是GRU单元隐藏层维度
vocab_inp_size = len(inp_lang.word_index)+1
# inp_lang.word_index能获取词到编码的映射关系字典，有多少个词就有多少个键值对，比如在只有两个词的时候，{'hello':1,'world':2},
# 对这个字典用len方法，就能获取一个有多少个词了
vocab_tar_size = len(targ_lang.word_index)+1

dataset = tf.data.Dataset.from_tensor_slices((input_tensor_train, target_tensor_train)).shuffle(BUFFER_SIZE)
dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)

看一眼数据维度：

example_input_batch, example_target_batch = next(iter(dataset))
example_input_batch.shape, example_target_batch.shape

# 输出：
# (TensorShape([64, 16]), TensorShape([64, 11]))

原文在这里会介绍下模型结构，带图，我这里就不贴出来了，自己看原文连接就行，当然，可能需要。
定义编码器：

class Encoder(tf.keras.Model):
   def __init__(self, vocab_size, embedding_dim, enc_units, batch_sz):
       super(Encoder, self).__init__()
       self.batch_sz = batch_sz
       self.enc_units = enc_units
       self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
       # 傻瓜式embedding操作，当然你也可以考虑使用训练好的embedding文件
       self.gru = tf.keras.layers.GRU(self.enc_units,
                                      return_sequences=True, 
                                      return_state=True,  
                                      recurrent_initializer='glorot_uniform')  # 循环核的初始化，Glorot均匀分布初始化方法
       # sequences，GRU的输出（即常见理论教材中的每个Cell的h），是GRU传给更上层GRU的
       # state，每个Cell的状态（即常见理论教材中的每个Cell的c），GRUCell传入下一个时间步的
       # 感觉这俩参数设置应该是继承LSTM模块而来，GRU是精简的LSTM，只有个h，或者说h和c相等，理论部分自行百度
       # 为了方便理解，后续会贴上示例
       
   def call(self, x, hidden):
       x = self.embedding(x)
       output, state = self.gru(x, initial_state = hidden)
       return output, state

   def initialize_hidden_state(self):
       return tf.zeros((self.batch_sz, self.enc_units))
"""    


"""

关于return_sequences和return_state的参数理解，
参考了https://zhuanlan.zhihu.com/p/85910281的文章，原文是对LSTM模块的这俩参数的设置，我改了改，重做了实验。

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input
from tensorflow.keras.layers import GRU
from numpy import array
tf.random.set_seed(2)
data = array([0.1, 0.2, 0.3]).reshape((1,3,1))
inputs1 = Input(shape=(3, 1))
gru1 = GRU(1)(inputs1)
model1 = Model(inputs=inputs1, outputs=gru1)
print(model1.predict(data),'\n')

默认都False，得到：[[-0.1320642]]
一个标量，模型输出

tf.random.set_seed(2)
gru2 = GRU(1, return_sequences=True)(inputs1)
model2 = Model(inputs=inputs1, outputs=gru2)
print(model2.predict(data),'\n')

return_sequences设为True，得到：
[[[-0.03509581],
[-0.08185742],
[-0.1320642 ]]]
三个时间步的输出，最后一步的cell输出也就是模型输出

tf.random.set_seed(2)
gru3 = GRU(1, return_state=True)(inputs1)
model3 = Model(inputs=inputs1, outputs=gru3)
print(model3.predict(data),'\n')

return_state设为True，得到：
[array([[-0.1320642]], dtype=float32), array([[-0.1320642]], dtype=float32)]
这个似乎有点费解，我估计是，首先得返回个模型输出，然后得返回最后一个cell的state，因为GRU的缘故，这俩恰好相等

tf.random.set_seed(2)
gru4 = GRU(1, return_sequences=True, return_state=True)(inputs1)
model4 = Model(inputs=inputs1, outputs=gru4)
print(model4.predict(data),'\n')

return_sequences和return_state同时设为True，得到两个值：
[array([[[-0.03509581],
[-0.08185742],
[-0.1320642 ]]], dtype=float32), array([[-0.1320642]], dtype=float32)]
一个模型输出，因为设置return_sequences为True，所以输出维度是1,3,1，
一个是最后一个cell的state，因为GRU的缘故，这个值和模型最后输出相等。

好，回到代码：

# 测试下编码器：
encoder = Encoder(vocab_inp_size, embedding_dim, units, BATCH_SIZE)

# 看一眼
sample_hidden = encoder.initialize_hidden_state()
sample_output, sample_hidden = encoder(example_input_batch, sample_hidden)
print ('Encoder output shape: (batch size, sequence length, units) {}'.format(sample_output.shape))
print ('Encoder Hidden state shape: (batch size, units) {}'.format(sample_hidden.shape))

# 输出：
# Encoder output shape: (batch size, sequence length, units) (64, 16, 1024)
# Encoder Hidden state shape: (batch size, units) (64, 1024)

注意力：

class BahdanauAttention(tf.keras.layers.Layer):
   def __init__(self, units):
       super(BahdanauAttention, self).__init__()
       self.W1 = tf.keras.layers.Dense(units)
       self.W2 = tf.keras.layers.Dense(units)
       self.V = tf.keras.layers.Dense(1)

   def call(self, query, values):
       hidden_with_time_axis = tf.expand_dims(query, 1)
       # 隐藏层的形状 == （批大小，隐藏层大小）
       # hidden_with_time_axis 的形状 == （批大小，1，隐藏层大小）
       # 这样做是为了执行加法以计算分数
       
       score = self.V(tf.nn.tanh(self.W1(values) + self.W2(hidden_with_time_axis)))
       # values就是Encoder最后一步输出，所以维度是（批大小，序列长度，隐藏层units）
       # score的形状 == （批大小，最大长度，1）
       # 我们在最后一个轴上得到 1， 因为我们把分数应用于 self.V
       # 在应用 self.V 之前，张量的形状是（批大小，最大长度，单位）

       attention_weights = tf.nn.softmax(score, axis=1)
       # 注意力权重 （attention_weights） 的形状 == （批大小，最大长度，1）

       context_vector = tf.reduce_sum(attention_weights * values, axis=1)
       # 上下文向量 （context_vector） 求和之后的形状 == （批大小，隐藏层大小）

       return context_vector, attention_weights

看一眼：

attention_layer = BahdanauAttention(10)
attention_result, attention_weights = attention_layer(sample_hidden, sample_output)

print("Attention result shape: (batch size, units) {}".format(attention_result.shape))
print("Attention weights shape: (batch_size, sequence_length, 1) {}".format(attention_weights.shape))

# 输出：
# Attention result shape: (batch size, units) (64, 1024)
# Attention weights shape: (batch_size, sequence_length, 1) (64, 16, 1)

解码器：

class Decoder(tf.keras.Model):
   def __init__(self, vocab_size, embedding_dim, dec_units, batch_sz):
       super(Decoder, self).__init__()
       self.batch_sz = batch_sz
       self.dec_units = dec_units
       self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
       self.gru = tf.keras.layers.GRU(self.dec_units,
                                      return_sequences=True,
                                      return_state=True,
                                      recurrent_initializer='glorot_uniform')
       self.fc = tf.keras.layers.Dense(vocab_size)
       # 和Encoder类差不多，就不过多注释了
       # 刚开始看到这里挺费解，这样直接用一个完整的GRU层，而不是GRUCell，解码时岂不是只有第一个起始符考虑了attention
       # 然后往下翻代码，发现是输出的每一步都会调用一次这个解码器

       self.attention = BahdanauAttention(self.dec_units)

   def call(self, x, hidden, enc_output):
       # 编码器输出 （enc_output） 的形状 == （批大小，最大长度，隐藏层大小）
       context_vector, attention_weights = self.attention(hidden, enc_output)

       # x 在通过嵌入层后的形状 == （批大小，1，嵌入维度）
       x = self.embedding(x)

       # x 在拼接 （concatenation） 后的形状 == （批大小，1，嵌入维度 + 隐藏层大小）
       x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)

       # 将合并后的向量传送到 GRU
       output, state = self.gru(x)

       # 输出的形状 == （批大小 * 1，隐藏层大小）
       output = tf.reshape(output, (-1, output.shape[2]))

       # 输出的形状 == （批大小，vocab）
       x = self.fc(output)
       
       return x, state, attention_weights

看一眼：

decoder = Decoder(vocab_tar_size, embedding_dim, units, BATCH_SIZE)

sample_decoder_output, _, _ = decoder(tf.random.uniform((64, 1)),
                                     sample_hidden, sample_output)

print ('Decoder output shape: (batch_size, vocab size) {}'.format(sample_decoder_output.shape))
# 输出：
# Decoder output shape: (batch_size, vocab size) (64, 4817)

定义优化器和损失函数：

optimizer = tf.keras.optimizers.Adam()
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')
# 损失函数带sparse，表示目标变量是未经one hot编码的分类变量
# from logits设为True表示网络输出未经softmax或者sigmoid函数处理，官方api手册表示这样计算结果更稳定
# 能一行一行看到这里的，这两点应该都懂

def loss_function(real, pred):
   mask = tf.math.logical_not(tf.math.equal(real, 0))
   # 这俩函数如其名，equal判断是否相等，返回T或者F，logical_not取反
   loss_ = loss_object(real, pred)
   # 传入损失函数
   mask = tf.cast(mask, dtype=loss_.dtype)
   # 变量类型转换，将mask转为loss_的数据类型
   # 一顿操作下来，如果real是0，mask就先是False，转张量后就变成0，如果real不是0，最后就变成1
   loss_ *= mask
   # mask机制，我个人理解就是，很多句子达不到指定长度，原句后面都是0的，就不计入loss，so easy
   
   return tf.reduce_mean(loss_)

检查点（基于对象保存）：

checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(optimizer=optimizer,
                                encoder=encoder,
                                decoder=decoder)

训练
1、将 输入 传送至 编码器，编码器返回 编码器输出 和 编码器隐藏层状态。
2、将编码器输出、编码器隐藏层状态和解码器输入（即 开始标记）传送至解码器。
3、解码器返回 预测 和 解码器隐藏层状态。
4、解码器隐藏层状态被传送回模型，预测被用于计算损失。
5、使用 教师强制 （teacher forcing） 决定解码器的下一个输入。
6、教师强制 是将 目标词 作为 下一个输入 传送至解码器的技术。
7、最后一步是计算梯度，并将其应用于优化器和反向传播。

@tf.function
# 这里用了个装饰器，tf就默认使用图计算，但是缺点就是，调试的时候不能即时查看变量值了，
# 比如对于一个tensor，用numpy()方法就会报错，如果想在函数计算过程中加个print来看看变量到底什么样的话，需要先把这个@tf.function注释掉
def train_step(inp, targ, enc_hidden):
   loss = 0

   with tf.GradientTape() as tape:
       enc_output, enc_hidden = encoder(inp, enc_hidden)
       # 先把初始输入和初始隐藏层状态传入解码器，然后得到解码器输出，顺便更新隐藏层状态，没啥好说的
       dec_hidden = enc_hidden
       # 经编码器计算后，隐藏层状态进入下一步，当作解码器的初始输入隐藏层状态

       dec_input = tf.expand_dims([targ_lang.word_index['']] * BATCH_SIZE, 1)
       # 维度变成batch size*1，解码器初始输入是句子起始符<start>，在下面的循环中，dec_input会不断更新

       for t in range(1, targ.shape[1]):
           # 注意这个for要在with tf.GradientTape() as tape包裹之内
           predictions, dec_hidden, _ = decoder(dec_input, dec_hidden, enc_output)
           # 将编码器输出 （enc_output） 传送至解码器

           loss += loss_function(targ[:, t], predictions)
           # loss累加

           dec_input = tf.expand_dims(targ[:, t], 1)
           # 使用教师强制，将目标词作为下一个输入

   batch_loss = (loss / int(targ.shape[1]))
   # 这个batch loss和我刚开始对字面理解的意思有点差异，
   # 看后面的代码可以知道，targ是一个batch的目标数据，维度是batch size*max targ length
   # 所以这个batch loss应该理解为一个batch数据上，平均每个时间步的loss
   # 另外查loss的代码可知，已经对batch取均值了

   variables = encoder.trainable_variables + decoder.trainable_variables
   # 获取待训练的变量
   
   gradients = tape.gradient(loss, variables)
   # 求导

   optimizer.apply_gradients(zip(gradients, variables))
   # 更新变量

   return batch_loss

EPOCHS = 10

for epoch in range(EPOCHS):
   start = time.time()

   enc_hidden = encoder.initialize_hidden_state()
   # 隐藏层状态初始化，用于第一次输入到编码器的，置零
   total_loss = 0

   for (batch, (inp, targ)) in enumerate(dataset.take(steps_per_epoch)):
       # 用enumerate打包之后再迭代，batch就是序号，这里也用来记录训练到第几个batch了
       # 封装好的数据集有个take方法，每次取steps_per_epoch条数据，
       # 其实不明白这里为什么不每次取batch size条数据，这样才符合逻辑
       batch_loss = train_step(inp, targ, enc_hidden)
       total_loss += batch_loss

       if batch % 100 == 0:
           print('Epoch {} Batch {} Loss {:.4f}'.format(epoch + 1,
                                                        batch,
                                                        batch_loss.numpy()))
   # 每 2 个周期（epoch），保存（检查点）一次模型
   if (epoch + 1) % 2 == 0:
       checkpoint.save(file_prefix = checkpoint_prefix)

   print('Epoch {} Loss {:.4f}'.format(epoch + 1,total_loss / steps_per_epoch))
   print('Time taken for 1 epoch {} sec\n'.format(time.time() - start))

翻译
1、评估函数类似于训练循环，不同之处在于在这里我们不使用 教师强制。每个时间步的解码器输入是其先前的预测、隐藏层状态和编码器输出。
2、当模型预测 结束标记 时停止预测。
3、存储 每个时间步的注意力权重。
4、请注意：对于一个输入，编码器输出仅计算一次。

def evaluate(sentence):
   attention_plot = np.zeros((max_length_targ, max_length_inp))
   # 一行就是输入序列的attention向量，维度就是输入序列的最大长度
   # 有max_length_targ行，就是每个时间步的target需要用一个attention向量

   sentence = preprocess_sentence(sentence)
   # 预处理，前面写好了的函数，忘了的话往上翻

   inputs = [inp_lang.word_index[i] for i in sentence.split(' ')]
   # 转编码
   inputs = tf.keras.preprocessing.sequence.pad_sequences([inputs],
                                                          maxlen=max_length_inp,
                                                          padding='post')
   # 填充or阶段
   
   inputs = tf.convert_to_tensor(inputs)
   # 转张量参与模型计算，但是我记忆中现在不需要这一步也能运行，实测也是如此

   result = ''

   hidden = [tf.zeros((1, units))]
   # 用来翻译的，所以第一个维度就是1了
   
   enc_out, enc_hidden = encoder(inputs, hidden)

   dec_hidden = enc_hidden
   dec_input = tf.expand_dims([targ_lang.word_index['']], 0)
   # tf计算时，最少二维，也就是传个向量进来，需要再包个列表符，有一定的tf经验的应该都懂

   for t in range(max_length_targ):
       predictions, dec_hidden, attention_weights = decoder(dec_input,
                                                            dec_hidden,
                                                            enc_out)
       attention_weights = tf.reshape(attention_weights, (-1, ))
       # 存储注意力权重以便后面制图
       # 因为输入只有一条，所以reshape的时候，第一个维度不加限制，设为-1就行，最后会变成16维的向量
       # 本例中，在reshape之前，attention weights的维度是(1,16,1)，可以参照前面的BahdanauAttention类注释来理解

       attention_plot[t] = attention_weights.numpy()
       # 一轮循环得到当前解码器时间步的词关于输入的attention向量，所以最后会得到个矩阵

       predicted_id = tf.argmax(predictions[0]).numpy()
       # 本例中，predictions维度是（1，4871），所以要取下标0
       # 用argmax获取最大概率词的下标（整个文档里一共4871个词）
       # 最后用numpy方法，把tensor转列表

       result += targ_lang.index_word[predicted_id] + ' '
       # 通过下标获取对应词

       if targ_lang.index_word[predicted_id] == '':
           return result, sentence, attention_plot
       # 如果预测结果为终止符，就停止循环，并返回结果

       dec_input = tf.expand_dims([predicted_id], 0)
       # 预测结果传入下一个时间步

   return result, sentence, attention_plot

画图函数，没啥好说的：

def plot_attention(attention, sentence, predicted_sentence):
   fig = plt.figure(figsize=(10,10))
   ax = fig.add_subplot(1, 1, 1)
   ax.matshow(attention, cmap='viridis')

   fontdict = {'fontsize': 14}

   ax.set_xticklabels([''] + sentence, fontdict=fontdict, rotation=90)
   ax.set_yticklabels([''] + predicted_sentence, fontdict=fontdict)

   ax.xaxis.set_major_locator(ticker.MultipleLocator(1))
   ax.yaxis.set_major_locator(ticker.MultipleLocator(1))

   plt.show()

翻译，并画注意力图，也没啥好说的：

def translate(sentence):
   result, sentence, attention_plot = evaluate(sentence)

   print('Input: %s' % (sentence))
   print('Predicted translation: {}'.format(result))

   attention_plot = attention_plot[:len(result.split(' ')), :len(sentence.split(' '))]
   plot_attention(attention_plot, sentence.split(' '), result.split(' '))

恢复检查点目录 （checkpoint_dir） 中最新的检查点：

checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))

见证奇迹的时刻：

translate(u'esta es mi vida.')

over