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='sparse_categorical_crossentropy')

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