RNN文本生成-TensorFlow2.0实现

• 加载数据，这里可以使用自己的数据集

# 加载数据
def get_data_from_file(train_file, batch_size, seq_size):
    with open(train_file, encoding='utf-8') as f:
        text = f.read()
        
    text = text.split()
    
    # 获取频繁词
    word_counts = Counter(text)
    # 按key降序排列
    sorted_vocab = sorted(word_counts, key=word_counts.get, reverse=True)
    # 创建词汇表
    int_to_vocab = {k: w for k, w in enumerate(sorted_vocab)}
    vocab_to_int = {w: k for k, w in int_to_vocab.items()}
    # 词汇表大小
    n_vocab = len(int_to_vocab)

    print('Vocabulary size', n_vocab)
    
    # 输入文本
    int_text = [vocab_to_int[w] for w in text]
    # 训练总需的批次数
    num_batches = int(len(int_text) / (seq_size * batch_size))
    in_text = int_text[:num_batches * batch_size * seq_size]
    
    # 后移
    out_text = np.zeros_like(in_text)
    out_text[:-1] = in_text[1:]
    out_text[-1] = in_text[0]
    
    in_text = np.reshape(in_text, (-1, seq_size))
    out_text = np.reshape(out_text, (-1, seq_size))
    
    return int_to_vocab, vocab_to_int, n_vocab, in_text, out_text

• 构建模型：由于这里只是原理的简单实现，因此只使用了一层单向的LSTM。在实际使用中，为了模型的效果可以使用多层双向的LSTM/GRU或是Transformer。

# 构建模型，这里使用的是单层的LSTM
class RNNModule(tf.keras.Model):
    def __init__(self, n_vocab, embedding_size, lstm_size):
        super(RNNModule, self).__init__()
        self.lstm_size = lstm_size
        
        # embedding shape:n_vocab * embedding_size
        self.embedding = tf.keras.layers.Embedding(n_vocab, embedding_size)
        self.lstm = tf.keras.layers.LSTM(lstm_size, return_state=True, return_sequences=True)
        self.dense = tf.keras.layers.Dense(n_vocab)
    
    # x shape: batch_size * seq_size
    def call(self, x, prev_state):
        embed = self.embedding(x)
        output, state_h, state_c = self.lstm(embed, prev_state)
        
        logits = self.dense(output)
        preds = tf.nn.softmax(logits)
        
        return logits, preds, (state_h, state_c)

    # 状态初始化
    def zero_state(self, batch_size):
        return [tf.zeros([batch_size, self.lstm_size]),
                tf.zeros([batch_size, self.lstm_size])]

• 训练函数

@tf.function
def train_func(inputs, targets, model, state, loss_func, optimizer):
  with tf.GradientTape() as tape:
      logits, _, state = model(inputs, state)

      loss = loss_func(targets, logits)

      gradients = tape.gradient(loss, model.trainable_variables)
      optimizer.apply_gradients(zip(gradients, model.trainable_variables))
      
      return loss

• 训练过程

# 训练
def train():
    int_to_vocab, vocab_to_int, n_vocab, in_text, out_text = get_data_from_file(args.train_file, args.batch_size, args.seq_size)
    # 文本长度
    len_data = in_text.shape[0]
    # 每个epoch训练的步数
    steps_per_epoch = len_data // args.batch_size
    
    # 创建dataset对象
    dataset = tf.data.Dataset.from_tensor_slices((in_text, out_text)).shuffle(args.buffer_size)
    dataset = dataset.batch(16, drop_remainder=True)
    
    # 加载模型、Encoder状态初始化
    model = RNNModule(n_vocab, args.embedding_size, args.lstm_size)
    state = model.zero_state(args.batch_size)
    
    # 选择优化器和损失项
    optimizer = tf.keras.optimizers.Adam()
    loss_func = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    
    # 开始
    for epoch in range(args.num_epochs):
        start = time.time()
        total_loss = list()
        state = model.zero_state(args.batch_size)

        for (batch, (inputs, targets)) in enumerate(dataset.take(steps_per_epoch)):
            
            
            # print ('Inputs shape: {}'.format(inputs.shape)) 
            # print ('targets shape: {}'.format(targets.shape))
            # Inputs shape: (16, 32)
            # targets shape: (16, 32)
            
            loss = train_func(inputs, targets, model, state, loss_func, optimizer)
            total_loss.append(loss)
            
            if batch % 100 == 0:
                print('Epoch: {}/{}'.format(epoch, args.num_epochs),
                      'Batch--> {}'.format(batch),
                      'Loss--> {}'.format(loss.numpy()))
                
            if batch % 300 == 0:
                predict(model, vocab_to_int, int_to_vocab, n_vocab)
                model.save_weights(checkpoint_prefix.format(epoch = epoch))
                
        plot_loss(total_loss)
        print ('Total time of this epoch is {}:'.format(time.time() - start))

完整实现代码

from __future__ import division
import argparse
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow import keras
import os
import time 
from collections import Counter

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# 超参数设置
def add_arguments(parser):
    parser.add_argument('--train_file', default='oliver.txt')
    parser.add_argument('--seq_size',type=int, default = 32,help='sequence length')
    parser.add_argument('--batch_size',type=int,default=16,help='batch size')
    parser.add_argument('--embedding_size',type=int,default=128,help='embedding hidden size')
    parser.add_argument('--lstm_size',type=int,default=128,help='lstm hidden size')
    parser.add_argument('--dropout_keep_prob',type=float,default=0.7,help='lstm dropout keep probability')
    parser.add_argument('--gradients_norm',type=int,default=5,help='norm to clip gradients')
    parser.add_argument('--initial_words',type=list,default=['I','am'],help='Initial words to start prediction from')
    parser.add_argument('--predict_top_k',type=int,default=5,help='top k results to sample word from')
    parser.add_argument('--num_epochs',type=int,default=20,help='number of epochs to train')
    parser.add_argument('--checkpoint_path',default='checkpoint',help='directory to store trained weights')
    parser.add_argument('--buffer_size',type=int,default=100,help='buffer size to use')
    

parser = argparse.ArgumentParser()
add_arguments(parser)
args = parser.parse_args()

# 加载数据
def get_data_from_file(train_file, batch_size, seq_size):
    with open(train_file, encoding='utf-8') as f:
        text = f.read()
        
    text = text.split()
    
    # 获取频繁词
    word_counts = Counter(text)
    # 按key降序排列
    sorted_vocab = sorted(word_counts, key=word_counts.get, reverse=True)
    # 创建词汇表
    int_to_vocab = {k: w for k, w in enumerate(sorted_vocab)}
    vocab_to_int = {w: k for k, w in int_to_vocab.items()}
    # 词汇表大小
    n_vocab = len(int_to_vocab)

    print('Vocabulary size', n_vocab)
    
    # 输入文本
    int_text = [vocab_to_int[w] for w in text]
    # 训练总需的批次数
    num_batches = int(len(int_text) / (seq_size * batch_size))
    in_text = int_text[:num_batches * batch_size * seq_size]
    
    # 后移
    out_text = np.zeros_like(in_text)
    out_text[:-1] = in_text[1:]
    out_text[-1] = in_text[0]
    
    in_text = np.reshape(in_text, (-1, seq_size))
    out_text = np.reshape(out_text, (-1, seq_size))
    
    return int_to_vocab, vocab_to_int, n_vocab, in_text, out_text

# 构建模型，这里使用的是单层的LSTM
class RNNModule(tf.keras.Model):
    def __init__(self, n_vocab, embedding_size, lstm_size):
        super(RNNModule, self).__init__()
        self.lstm_size = lstm_size
        
        # embedding shape:n_vocab * embedding_size
        self.embedding = tf.keras.layers.Embedding(n_vocab, embedding_size)
        self.lstm = tf.keras.layers.LSTM(lstm_size, return_state=True, return_sequences=True)
        self.dense = tf.keras.layers.Dense(n_vocab)
    
    # x shape: batch_size * seq_size
    def call(self, x, prev_state):
        embed = self.embedding(x)
        output, state_h, state_c = self.lstm(embed, prev_state)
        
        logits = self.dense(output)
        preds = tf.nn.softmax(logits)
        
        return logits, preds, (state_h, state_c)

    # 状态初始化
    def zero_state(self, batch_size):
        return [tf.zeros([batch_size, self.lstm_size]),
                tf.zeros([batch_size, self.lstm_size])]

# chekpoint设置
checkpoint_dir = './training_chekpoints'
checkpoint_prefix = os.path.join(checkpoint_dir,"ckpt_{epoch}")
checkpoint_callback = keras.callbacks.ModelCheckpoint(filepath = checkpoint_prefix,save_weights_only = True)

# 预测
def predict(model, vocab_to_int, int_to_vocab, n_vocab):
    # 根据概率选择新的词
    def get_word(int_pred, n_vocab):
        p = np.squeeze(int_pred)
        p[p.argsort()][:-5] = 0
        p = p / np.sum(p)
        word = np.random.choice(n_vocab, 1, p=p)[0]

        return word

    val_state = model.zero_state(1)
    words = args.initial_words
    # 生成文本
    for word in words:
        int_word = tf.convert_to_tensor(
            [[vocab_to_int[word]]], dtype=tf.float32)
        _, int_pred, val_state = model(int_word, val_state)
    int_pred = int_pred.numpy()
    int_word = get_word(int_pred, n_vocab)
    words.append(int_to_vocab[int_word])
    # 保证生成的文本长度为50
    for _ in range(50):
        int_word = tf.convert_to_tensor(
            [[int_word]], dtype=tf.float32)
        _, int_pred, val_state = model(int_word, val_state)
        int_pred = int_pred.numpy()
        int_word = get_word(int_pred, n_vocab)
        words.append(int_to_vocab[int_word])
        
    print('Generated sentences are: {}'.format(' '.join(words)))


@tf.function
def train_func(inputs, targets, model, state, loss_func, optimizer):
  with tf.GradientTape() as tape:
      logits, _, state = model(inputs, state)

      loss = loss_func(targets, logits)

      gradients = tape.gradient(loss, model.trainable_variables)
      optimizer.apply_gradients(zip(gradients, model.trainable_variables))
      
      return loss


# 训练
def train():
    int_to_vocab, vocab_to_int, n_vocab, in_text, out_text = get_data_from_file(args.train_file, args.batch_size, args.seq_size)
    # 文本长度
    len_data = in_text.shape[0]
    # 每个epoch训练的步数
    steps_per_epoch = len_data // args.batch_size
    
    # 创建dataset对象
    dataset = tf.data.Dataset.from_tensor_slices((in_text, out_text)).shuffle(args.buffer_size)
    dataset = dataset.batch(16, drop_remainder=True)
    
    # 加载模型、Encoder状态初始化
    model = RNNModule(n_vocab, args.embedding_size, args.lstm_size)
    state = model.zero_state(args.batch_size)
    
    # 选择优化器和损失项
    optimizer = tf.keras.optimizers.Adam()
    loss_func = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    
    # 开始
    for epoch in range(args.num_epochs):
        start = time.time()
        total_loss = list()
        state = model.zero_state(args.batch_size)

        for (batch, (inputs, targets)) in enumerate(dataset.take(steps_per_epoch)):
            
            
            # print ('Inputs shape: {}'.format(inputs.shape)) 
            # print ('targets shape: {}'.format(targets.shape))
            # Inputs shape: (16, 32)
            # targets shape: (16, 32)
            
            loss = train_func(inputs, targets, model, state, loss_func, optimizer)
            total_loss.append(loss)
            
            if batch % 100 == 0:
                print('Epoch: {}/{}'.format(epoch, args.num_epochs),
                      'Batch--> {}'.format(batch),
                      'Loss--> {}'.format(loss.numpy()))
                
            if batch % 300 == 0:
                predict(model, vocab_to_int, int_to_vocab, n_vocab)
                model.save_weights(checkpoint_prefix.format(epoch = epoch))
                
        plot_loss(total_loss)
        print ('Total time of this epoch is {}:'.format(time.time() - start))
        
# 绘制损失变化图
def plot_loss(losses):
    plt.plot(losses)
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.show()
    
if __name__ == '__main__':
    train()