TensorFlow文本生成（AI 写诗）

时间：04-17来源：作者：点击数：42

俗话说的好，“熟读唐诗三百首，不会作诗也会吟”。吟诗作对我是做不到的，那就训练一个模型，让它去“背书”吧，背完了再看看它学的怎么样。

当然这里的写诗肯定不会照搬已存在诗句，而是根据它对看过的诗的理解，再给它一两个字或者一句诗，甚至一个字也不给，然后写出格式正确的诗。当然这里的格式正确不代表能有明确的意思，大概率是句子看着像那么回事，但整句诗没什么逻辑，这与数据量的大小、训练次数、模型结构有关。

本次使用的是 TensorFlow2，采用基于循环神经网络的 GRU ，所有的代码会放在 GitHub 里面：https://github.com/Stevengz/Poem_compose

需要用到的库：

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import losses
import numpy as np
import os

数据预处理

首先得获取诗的数据，很多网站都保存了相当数量的诗句，本来准备随便找一个自己爬的，结果发现了一位大神维护了一个诗词仓库：https://github.com/chinese-poetry/chinese-poetry，唐诗宋词都有。里面有直接用 json 格式保存的内容，选择一些把诗句内容提取出来就行了，本来想自己提取的，结果又有人直接把其中一个文件内的诗句提取并转换为数字形式了，大概千首诗，不过也提供了字符和数字之间的相互索引，全部存在了一个压缩文件里面，我会把这个文件也放在仓库里面，可以自行提取。

这样数据集就解决了，给我省了不少事。接下来就把数据集文件和内容介绍一下，压缩文件里面有三个 npy 文件，一个保存诗词内容，另外两个是字符和数字分别相互索引的字典。

把汉字转换为数字是为了便于模型计算，不同的数字对应于不同的汉字，只要能够将不同的字区分开，数字或汉字在计算机眼里是一样的。

看一下它们长什么样子，data 里面每首诗的内容都存在一个列表里面，随便挑一首看看：

把它们转换为汉字：

可以看到这一首诗里面最多的是</s>，它代表空格，因为每首诗的长度是不一样的，模型训练是要确保所有输入的长度相同（为 125），不够长的位置就用空格填充。而<START>和<EOP>表示诗句开始和结束，也是为了帮助模型能够分辨。

接下来就开始导入数据：

# 导入数据集
data = np.load('data.npy', allow_pickle=True).tolist()
# 将二维数据变为一维，就是将所有诗的每个字放在同一列表内
data_line = np.array([word for poem in data for word in poem])
# 两个字典索引
ix2word = np.load('ix2word.npy', allow_pickle=True).item()
word2ix = np.load('word2ix.npy', allow_pickle=True).item()

上面将诗句展开了，接下来使用 tensorflow 将每首诗提取出来，再将一定数量的诗分为一批，将所有诗分批隔开，后面训练时将数据一批批导入

# 每批大小
BATCH_SIZE = 64
# 缓冲区大小
BUFFER_SIZE = 10000
# 诗的长度
poem_size = 125
# 创建训练样本
poem_dataset = tf.data.Dataset.from_tensor_slices(data_line)
# 将每首诗提取出来
poems = poem_dataset.batch(poem_size + 1, drop_remainder=True)
# 切分成输入和输出，一个去头一个去尾，输入输出的长度都是 125
def split_input_target(chunk):
    input_text = chunk[:-1]
    target_text = chunk[1:]
    return input_text, target_text
dataset = poems.map(split_input_target)
# 分批并随机打乱
dataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)

这里所谓输入和输出的长度，是训练时确定的数据结构，两者长度不一定要相等。这不影响模型训练完后的数据输入和输出的格式，此时即使是一个字也可以输入，每次输出数量为批次大小的字，要获取一定数量的句子则一步步生成。

模型构建

这里就简单地使用三层结构来定义模型：

def GRU_model(vocab_size, embedding_dim, units, batch_size):
    model = keras.Sequential([
        layers.Embedding(vocab_size,
                         embedding_dim,
                         batch_input_shape=[batch_size, None]),
        layers.GRU(units,
                   return_sequences=True,
                   stateful=True,
                   recurrent_initializer='glorot_uniform'),
        layers.Dense(vocab_size)
    ])
    return model

其中vocab_size是词集的长度，即训练集中有多少个不同的字符。

三层结构中：Embedding是输入层，一个可训练的对照表，它会将每个字符的数字映射到一个 embedding_dim 维度的向量。；GRU是一种 RNN 类型，其大小由 units 指定；Dense是输出层，带有 vocab_size 个输出。

实例化构建：

# 嵌入的维度
embedding_dim = 64
# RNN 的单元数量
units = 128

# 创建模型
model = GRU_model(vocab_size=len(ix2word),
                  embedding_dim=embedding_dim,
                  units=units,
                  batch_size=BATCH_SIZE)
model.summary()

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

模型训练完成后，不管输入多少个字，都会生成一个 (batch_size, vocab_size) 大小的结果，其中每个 (1, vocab_size) 表示词集中每个字的概率，根据概率选择一个字为输出，而 batch_size 则是每次生成字的数量。

训练模型

先添加优化器和损失函数：

# 损失函数
def loss(labels, logits):
    return tf.keras.losses.sparse_categorical_crossentropy(labels,
                                                           logits,
                                                           from_logits=True)
model.compile(optimizer='adam', loss=loss)

配置检查点后训练：

# 检查点目录
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)

# 训练周期
EPOCHS = 20
history = model.fit(dataset, epochs=EPOCHS, callbacks=[checkpoint_callback])

检查点的作用是每个周期训练完后将模型进行保存

写诗

将批大小设定为 1，使过程简单一点，RNN 状态从时间步传递到时间步的方式，模型建立好之后只接受固定的批大小。

model = GRU_model(len(ix2word), embedding_dim, units=units, batch_size=1)
# 从检查点中恢复权重
model.load_weights(tf.train.latest_checkpoint(checkpoint_dir))
model.build(tf.TensorShape([1, None]))

接下来就是生成文本了

首先设置起始字符串，初始化 RNN 状态并设置要生成的字符个数。用起始字符串和 RNN 状态，获取下一个字符的预测分布。然后，用分类分布计算预测字符的索引。

把这个预测字符当作模型的下一个输入。模型返回的 RNN 状态被输送回模型。现在，模型有更多上下文可以学习，而非只有一个字。在预测出下一个字后，更改过的 RNN 状态被再次输送回模型。模型就是这样，通过不断从前面预测的字符获得更多上下文，进行学习。

def generate_text(model, start_string):
    # 要生成的字符个数
    num_generate = 120
    # 空字符串用于存储结果
    poem_generated = []
    temperature = 1.0
    # 将整个输入直接导入
    input_eval = [word2ix[s] for s in prefix_words + start_string]
    # 添加开始标识
    input_eval.insert(0, word2ix['<START>'])
    input_eval = tf.expand_dims(input_eval, 0)
    model.reset_states()
    for i in range(num_generate):
        predictions = model(input_eval)
        # 删除批次的维度
        predictions = tf.squeeze(predictions, 0)
        # 用分类分布预测模型返回的字符
        predictions = predictions / temperature
        predicted_id = tf.random.categorical(predictions,
                                             num_samples=1)[-1, 0].numpy()
        # 把预测字符和前面的隐藏状态一起传递给模型作为下一个输入
        input_eval = tf.expand_dims([predicted_id], 0)
        poem_generated.append(ix2word[predicted_id])
    # 删除多余的字
    del poem_generated[poem_generated.index('<EOP>'):]
    return (start_string + ''.join(poem_generated))

print(generate_text(model, start_string="水"))