Tensorflow2.0之文本生成莎士比亚作品

我们将使用 Andrej Karpathy 在《循环神经网络不合理的有效性》一文中提供的莎士比亚作品数据集。给定此数据中的一个字符序列 （“Shakespear”），训练一个模型以预测该序列的下一个字符（“e”）。通过重复调用该模型，可以生成更长的文本序列。

1、导入数据

请参考Tensorflow2.0加载和预处理数据的方法汇总中的第八部分：导入文本（用于文本生成）。

2、创建模型

使用 tf.keras.Sequential 定义模型。在这个例子中，我们使用了三个层来定义模型：

• tf.keras.layers.Embedding：输入层。一个可训练的对照表，它会将每个字符的数字映射到一个 embedding_dim 维度的向量。
• tf.keras.layers.GRU：一种 RNN 的类型，其大小由 units=rnn_units 指定。
• tf.keras.layers.Dense：输出层，带有 vocab_size 个输出。

vocab_size = len(vocab)  # 词集的长度
embedding_dim = 256  # 嵌入的维度
rnn_units = 1024  # RNN 的单元数量

def build_model(vocab_size, embedding_dim, rnn_units, batch_size):
  model = tf.keras.Sequential([
    tf.keras.layers.Embedding(vocab_size, embedding_dim,
                              batch_input_shape=[batch_size, None]),
    tf.keras.layers.GRU(rnn_units,
                        return_sequences=True,
                        stateful=True,
                        recurrent_initializer='glorot_uniform'),
    tf.keras.layers.Dense(vocab_size)
  ])
  return model
model = build_model(vocab_size=len(vocab),
  					embedding_dim=embedding_dim,
  					rnn_units=rnn_units,
  					batch_size=BATCH_SIZE)
model.summary()

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
embedding (Embedding)        (64, None, 256)           16640     
_________________________________________________________________
gru (GRU)                    (64, None, 1024)          3938304   
_________________________________________________________________
dense (Dense)                (64, None, 65)            66625     
=================================================================
Total params: 4,021,569
Trainable params: 4,021,569
Non-trainable params: 0
_________________________________________________________________

对于每个字符，模型会查找嵌入，把嵌入当作输入运行 GRU 一个时间步，并用 Dense 层生成逻辑回归 ，预测下一个字符的对数可能性。

3、训练

3.1 编译模型

model.compile(optimizer='adam',
              loss=tf.keras.losses.sparse_categorical_crossentrop(from_logits=True))

3.2 配置检查点

# 检查点保存至的目录
checkpoint_dir = './training_checkpoints'

# 检查点的文件名
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")

checkpoint_callback=tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_prefix,
    save_weights_only=True)

3.3 训练模型

history = model.fit(dataset, epochs=examples_per_epoch, callbacks=[checkpoint_callback])

4、预测

下面的代码块用来生成文本：

• 首先设置起始字符串，初始化 RNN 状态并设置要生成的字符个数。
• 用起始字符串和 RNN 状态，获取下一个字符的预测分布。
• 然后，用分类分布计算预测字符的索引。把这个预测字符当作模型的下一个输入。
• 模型返回的 RNN 状态被输送回模型。现在，模型有更多上下文可以学习，而非只有一个字符。在预测出下一个字符后，更改过的 RNN 状态被再次输送回模型。模型就是这样，通过不断从前面预测的字符获得更多上下文，进行学习。
  

如上图所示，这里我们希望实现的功能是输入一个样本，设有n个字符，模型将输出每个输入字符后的一个字符，即共输出n个字符，然后将新得到的n个字符作为输入，再次输入模型，共重复这个步骤指定次数。

由于设置 GRU 隐藏状态的时候必须指定批次大小，所以模型建立好之后只能接受固定的批次大小。

若要使用不同的 batch_size 来运行模型，我们需要重建模型并从检查点中恢复权重。

4.1 重建模型

tf.train.latest_checkpoint(checkpoint_dir)
model = build_model(vocab_size, embedding_dim, rnn_units, batch_size=1)

model.load_weights(tf.train.latest_checkpoint(checkpoint_dir))

model.build(tf.TensorShape([1, None]))
model.summary()

Model: "sequential_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
embedding_1 (Embedding)      (1, None, 256)            16640     
_________________________________________________________________
gru_1 (GRU)                  (1, None, 1024)           3938304   
_________________________________________________________________
dense_1 (Dense)              (1, None, 65)             66625     
=================================================================
Total params: 4,021,569
Trainable params: 4,021,569
Non-trainable params: 0
_________________________________________________________________

4.2 生成文本

def generate_text(model, start_string):
  # 评估步骤（用学习过的模型生成文本）

  # 要生成的字符个数
  num_generate = 1000

  # 将起始字符串转换为数字（向量化）
  input_eval = [char2idx[s] for s in start_string]
  input_eval = tf.expand_dims(input_eval, 0)

  # 空字符串用于存储结果
  text_generated = []

  # 这里批大小为 1
  model.reset_states()
  for i in range(num_generate):
      predictions = model(input_eval)
      # 删除批次的维度
      predictions = tf.squeeze(predictions, 0)

      # 用分类分布预测模型返回的字符
      predicted_id = tf.random.categorical(predictions, num_samples=1)[-1,0].numpy()

      # 把预测字符和前面的隐藏状态一起传递给模型作为下一个输入
      input_eval = tf.expand_dims([predicted_id], 0)

      text_generated.append(idx2char[predicted_id])

  return (start_string + ''.join(text_generated))

print(generate_text(model, start_string=u"ROMEO: "))

在这里输入了7个字符，那么每次模型输出的形状为 (1, 7, 65)，我们只需要后面两个维度 (7, 65)，所以使用 tf.squeeze 函数去掉第一个维度；使用 tf.random.categorical 函数来确定每一行中最大概率所对应的索引，然后将此索引（再增加第一维度后）重新作为输入，重复以上步骤直到预测完指定的1000个字符为止。

最终得到结果：

ROMEO: it may be see, I say.
Elong where I have sea loved for such heart
As of all desperate in your colls?
On how much to purwed esumptrues as we,
But taker appearing our great Isabel,;
Of your brother's needs.
I cannot but one hour, by nimwo and ribs
After 't? O Pedur, break our manory,
The shadot bestering eyes write; onfility;
Indeed I am possips And feated with others and throw it?

CAPULET:
O, not the ut with mine own sort.
But, with your souls, sir, well we would he,
And videwith the sungesoy begins, revell;
Much it in secart.

PROSPERO:
Villain, I darry record;
In sea--lodies, nor that I do I were stir,
You appointed with that sed their o tailor and hope left fear'd,
I so; that your looks stand up,
Comes I truly see this last weok not the
sul us.

CAMILLO:
You did and ever sea,
Into these hours: awake! Ro with mine enemies,
Were werx'd in everlawacted man been to alter
As Lewis could smile to his.

Farthus:
Marry! I'll do lose a man see me
To no drinking often hat back on an illing mo