1、目录结构
2、入口类
# coding = utf-8 """ 注意：RNN使用的数据为序列化的数据 RNN网络：主要由多个LSTM计算单元组成，依靠BPTT算法进行时序展开 LSTM:含有保留门和忘记门，是一个多输入多输出的网络结构。 LSTM具备抑制梯度特性 """ # import numpy as np # import tensorflow as tf # from .models.model import rnn_model # from .dataset.poems import process_poems,generate_batch import argparse import sys sys.path.append(r'D:\study\python-数据分析\深度学习\RNN网络\inference') def parse_args(): """ 参数设定 :return: """ #参数描述 parser = argparse.ArgumentParser(description='Intelligence Poem and Lyric Writer.') help_ = 'you can set this value in terminal --write value can be poem or lyric.' parser.add_argument('-w', '--write', default='poem', choices=['poem', 'lyric'], help=help_) help_ = 'choose to train or generate.' #训练 parser.add_argument('--train', dest='train', action='store_true', help=help_) #测试 parser.add_argument('--no-train', dest='train', action='store_false', help=help_) parser.set_defaults(train=False) args_ = parser.parse_args() return args_ if __name__ == '__main__': args = parse_args() if args.write == 'poem': from inference import tang_poems if args.train: tang_poems.main(True) #训练 else: tang_poems.main(False) #测试 elif args.write == 'lyric': from inference import song_lyrics print(args.train) if args.train: song_lyrics.main(True) else: song_lyrics.main(False) else: print('[INFO] write option can only be poem or lyric right now.')
3、tang_poems.py
# -*- coding: utf-8 -*- # file: tang_poems.py import collections import os import sys import numpy as np import tensorflow as tf from models.model import rnn_model from dataset.poems import process_poems, generate_batch import heapq tf.app.flags.DEFINE_integer('batch_size', 64, 'batch size.') tf.app.flags.DEFINE_float('learning_rate', 0.01, 'learning rate.') # set this to 'main.py' relative path tf.app.flags.DEFINE_string('checkpoints_dir', os.path.abspath('./checkpoints/poems/'), 'checkpoints save path.') tf.app.flags.DEFINE_string('file_path', os.path.abspath('./dataset/data/poems.txt'), 'file name of poems.') tf.app.flags.DEFINE_string('model_prefix', 'poems', 'model save prefix.') tf.app.flags.DEFINE_integer('epochs', 50, 'train how many epochs.') FLAGS = tf.app.flags.FLAGS start_token = 'G' end_token = 'E' def run_training(): #模型保存路径配置 if not os.path.exists(os.path.dirname(FLAGS.checkpoints_dir)): os.mkdir(os.path.dirname(FLAGS.checkpoints_dir)) if not os.path.exists(FLAGS.checkpoints_dir): os.mkdir(FLAGS.checkpoints_dir) #1、诗集数据处理 poems_vector, word_to_int, vocabularies = process_poems(FLAGS.file_path) #2、生成批量数据用于训练 batches_inputs, batches_outputs = generate_batch(FLAGS.batch_size, poems_vector, word_to_int) input_data = tf.placeholder(tf.int32, [FLAGS.batch_size, None]) output_targets = tf.placeholder(tf.int32, [FLAGS.batch_size, None]) #3、建立模型 end_points = rnn_model(model='lstm', input_data=input_data, output_data=output_targets, vocab_size=len( vocabularies), rnn_size=128, num_layers=2, batch_size=64, learning_rate=FLAGS.learning_rate) saver = tf.train.Saver(tf.global_variables()) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) #4、开始训练 with tf.Session() as sess: # sess = tf_debug.LocalCLIDebugWrapperSession(sess=sess) # sess.add_tensor_filter("has_inf_or_nan", tf_debug.has_inf_or_nan) sess.run(init_op) start_epoch = 0 checkpoint = tf.train.latest_checkpoint(FLAGS.checkpoints_dir) if checkpoint: saver.restore(sess, checkpoint) print("[INFO] restore from the checkpoint {0}".format(checkpoint)) start_epoch += int(checkpoint.split('-')[-1]) print('[INFO] start training...') try: for epoch in range(start_epoch, FLAGS.epochs): n = 0 n_chunk = len(poems_vector) // FLAGS.batch_size for batch in range(n_chunk): loss, _, _ = sess.run([ end_points['total_loss'], end_points['last_state'], end_points['train_op'] ], feed_dict={input_data: batches_inputs[n], output_targets: batches_outputs[n]}) n += 1 print('[INFO] Epoch: %d , batch: %d , training loss: %.6f' % (epoch, batch, loss)) if epoch % 6 == 0: saver.save(sess, './model/', global_step=epoch) #saver.save(sess, os.path.join(FLAGS.checkpoints_dir, FLAGS.model_prefix), global_step=epoch) except KeyboardInterrupt: print('[INFO] Interrupt manually, try saving checkpoint for now...') saver.save(sess, os.path.join(FLAGS.checkpoints_dir, FLAGS.model_prefix), global_step=epoch) print('[INFO] Last epoch were saved, next time will start from epoch {}.'.format(epoch)) def to_word(predict, vocabs): t = np.cumsum(predict) s = np.sum(predict) sample = int(np.searchsorted(t, np.random.rand(1) * s)) if sample > len(vocabs): sample = len(vocabs) - 1 return vocabs[sample] def gen_poem(begin_word): batch_size = 1 print('[INFO] loading corpus from %s' % FLAGS.file_path) poems_vector, word_int_map, vocabularies = process_poems(FLAGS.file_path) input_data = tf.placeholder(tf.int32, [batch_size, None]) end_points = rnn_model(model='lstm', input_data=input_data, output_data=None, vocab_size=len( vocabularies), rnn_size=128, num_layers=2, batch_size=64, learning_rate=FLAGS.learning_rate) saver = tf.train.Saver(tf.global_variables()) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) with tf.Session() as sess: sess.run(init_op) #checkpoint = tf.train.latest_checkpoint(FLAGS.checkpoints_dir) checkpoint = tf.train.latest_checkpoint('./model/') #saver.restore(sess, checkpoint) saver.restore(sess, './model/-24') x = np.array([list(map(word_int_map.get, start_token))]) [predict, last_state] = sess.run([end_points['prediction'], end_points['last_state']], feed_dict={input_data: x}) if begin_word: word = begin_word else: word = to_word(predict, vocabularies) poem = '' while word != end_token: print ('runing') poem += word x = np.zeros((1, 1)) x[0, 0] = word_int_map[word] [predict, last_state] = sess.run([end_points['prediction'], end_points['last_state']], feed_dict={input_data: x, end_points['initial_state']: last_state}) word = to_word(predict, vocabularies) # word = words[np.argmax(probs_)] return poem def pretty_print_poem(poem): poem_sentences = poem.split('。') for s in poem_sentences: if s != '' and len(s) > 10: print(s + '。') def main(is_train): if is_train: print('[INFO] train tang poem...') run_training() else: print('[INFO] write tang poem...') begin_word = input('输入起始字:') #begin_word = '我' poem2 = gen_poem(begin_word) pretty_print_poem(poem2) if __name__ == '__main__': tf.app.run()