Não pode escolher mais do que 25 tópicos
Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.
234 linhas
11 KiB
234 linhas
11 KiB
import tensorflow as tf
|
|
import seq2seq
|
|
import bleu
|
|
import reader
|
|
from os import path
|
|
import random
|
|
|
|
|
|
class Model():
|
|
|
|
def __init__(self, train_input_file, train_target_file,
|
|
test_input_file, test_target_file, vocab_file,
|
|
num_units, layers, dropout,
|
|
batch_size, learning_rate, output_dir,
|
|
save_step = 100, eval_step = 1000,
|
|
param_histogram=False, restore_model=False,
|
|
init_train=True, init_infer=False):
|
|
self.num_units = num_units
|
|
self.layers = layers
|
|
self.dropout = dropout
|
|
self.batch_size = batch_size
|
|
self.learning_rate = learning_rate
|
|
self.save_step = save_step
|
|
self.eval_step = eval_step
|
|
self.param_histogram = param_histogram
|
|
self.restore_model = restore_model
|
|
self.init_train = init_train
|
|
self.init_infer = init_infer
|
|
|
|
if init_train:
|
|
self.train_reader = reader.SeqReader(train_input_file,
|
|
train_target_file, vocab_file, batch_size)
|
|
self.train_reader.start()
|
|
self.train_data = self.train_reader.read()
|
|
self.eval_reader = reader.SeqReader(test_input_file, test_target_file,
|
|
vocab_file, batch_size)
|
|
self.eval_reader.start()
|
|
self.eval_data = self.eval_reader.read()
|
|
|
|
self.model_file = path.join(output_dir, 'model.ckpl')
|
|
self.log_writter = tf.summary.FileWriter(output_dir)
|
|
|
|
if init_train:
|
|
self._init_train()
|
|
self._init_eval()
|
|
|
|
if init_infer:
|
|
self.infer_vocabs = reader.read_vocab(vocab_file)
|
|
self.infer_vocab_indices = dict((c, i) for i, c in
|
|
enumerate(self.infer_vocabs))
|
|
self._init_infer()
|
|
self.reload_infer_model()
|
|
|
|
|
|
def gpu_session_config(self):
|
|
config = tf.ConfigProto()
|
|
config.gpu_options.allow_growth = True
|
|
return config
|
|
|
|
|
|
def _init_train(self):
|
|
self.train_graph = tf.Graph()
|
|
with self.train_graph.as_default():
|
|
self.train_in_seq = tf.placeholder(tf.int32, shape=[self.batch_size, None])
|
|
self.train_in_seq_len = tf.placeholder(tf.int32, shape=[self.batch_size])
|
|
self.train_target_seq = tf.placeholder(tf.int32, shape=[self.batch_size, None])
|
|
self.train_target_seq_len = tf.placeholder(tf.int32, shape=[self.batch_size])
|
|
output = seq2seq.seq2seq(self.train_in_seq, self.train_in_seq_len,
|
|
self.train_target_seq, self.train_target_seq_len,
|
|
len(self.train_reader.vocabs),
|
|
self.num_units, self.layers, self.dropout)
|
|
self.train_output = tf.argmax(tf.nn.softmax(output), 2)
|
|
self.loss = seq2seq.seq_loss(output, self.train_target_seq,
|
|
self.train_target_seq_len)
|
|
params = tf.trainable_variables()
|
|
gradients = tf.gradients(self.loss, params)
|
|
clipped_gradients, _ = tf.clip_by_global_norm(
|
|
gradients, 0.5)
|
|
self.train_op = tf.train.AdamOptimizer(
|
|
learning_rate=self.learning_rate
|
|
).apply_gradients(zip(clipped_gradients,params))
|
|
if self.param_histogram:
|
|
for v in tf.trainable_variables():
|
|
tf.summary.histogram('train_' + v.name, v)
|
|
tf.summary.scalar('loss', self.loss)
|
|
self.train_summary = tf.summary.merge_all()
|
|
self.train_init = tf.global_variables_initializer()
|
|
self.train_saver = tf.train.Saver()
|
|
self.train_session = tf.Session(graph=self.train_graph,
|
|
config=self.gpu_session_config())
|
|
|
|
|
|
def _init_eval(self):
|
|
self.eval_graph = tf.Graph()
|
|
with self.eval_graph.as_default():
|
|
self.eval_in_seq = tf.placeholder(tf.int32, shape=[self.batch_size, None])
|
|
self.eval_in_seq_len = tf.placeholder(tf.int32, shape=[self.batch_size])
|
|
self.eval_output = seq2seq.seq2seq(self.eval_in_seq,
|
|
self.eval_in_seq_len, None, None,
|
|
len(self.eval_reader.vocabs),
|
|
self.num_units, self.layers, self.dropout)
|
|
if self.param_histogram:
|
|
for v in tf.trainable_variables():
|
|
tf.summary.histogram('eval_' + v.name, v)
|
|
self.eval_summary = tf.summary.merge_all()
|
|
self.eval_saver = tf.train.Saver()
|
|
self.eval_session = tf.Session(graph=self.eval_graph,
|
|
config=self.gpu_session_config())
|
|
|
|
|
|
def _init_infer(self):
|
|
self.infer_graph = tf.Graph()
|
|
with self.infer_graph.as_default():
|
|
self.infer_in_seq = tf.placeholder(tf.int32, shape=[1, None])
|
|
self.infer_in_seq_len = tf.placeholder(tf.int32, shape=[1])
|
|
self.infer_output = seq2seq.seq2seq(self.infer_in_seq,
|
|
self.infer_in_seq_len, None, None,
|
|
len(self.infer_vocabs),
|
|
self.num_units, self.layers, self.dropout)
|
|
self.infer_saver = tf.train.Saver()
|
|
self.infer_session = tf.Session(graph=self.infer_graph,
|
|
config=self.gpu_session_config())
|
|
|
|
|
|
|
|
def train(self, epochs, start=0):
|
|
if not self.init_train:
|
|
raise Exception('Train graph is not inited!')
|
|
with self.train_graph.as_default():
|
|
if path.isfile(self.model_file + '.meta') and self.restore_model:
|
|
print("Reloading model file before training.")
|
|
self.train_saver.restore(self.train_session, self.model_file)
|
|
else:
|
|
self.train_session.run(self.train_init)
|
|
total_loss = 0
|
|
for step in range(start, epochs):
|
|
data = next(self.train_data)
|
|
in_seq = data['in_seq']
|
|
in_seq_len = data['in_seq_len']
|
|
target_seq = data['target_seq']
|
|
target_seq_len = data['target_seq_len']
|
|
output, loss, train, summary = self.train_session.run(
|
|
[self.train_output, self.loss, self.train_op, self.train_summary],
|
|
feed_dict={
|
|
self.train_in_seq: in_seq,
|
|
self.train_in_seq_len: in_seq_len,
|
|
self.train_target_seq: target_seq,
|
|
self.train_target_seq_len: target_seq_len})
|
|
total_loss += loss
|
|
self.log_writter.add_summary(summary, step)
|
|
if step % self.save_step == 0:
|
|
self.train_saver.save(self.train_session, self.model_file)
|
|
print("Saving model. Step: %d, loss: %f" % (step,
|
|
total_loss / self.save_step))
|
|
# print sample output
|
|
sid = random.randint(0, self.batch_size-1)
|
|
input_text = reader.decode_text(in_seq[sid],
|
|
self.eval_reader.vocabs)
|
|
output_text = reader.decode_text(output[sid],
|
|
self.train_reader.vocabs)
|
|
target_text = reader.decode_text(target_seq[sid],
|
|
self.train_reader.vocabs).split(' ')[1:]
|
|
target_text = ' '.join(target_text)
|
|
print('******************************')
|
|
print('src: ' + input_text)
|
|
print('output: ' + output_text)
|
|
print('target: ' + target_text)
|
|
if step % self.eval_step == 0:
|
|
bleu_score = self.eval(step)
|
|
print("Evaluate model. Step: %d, score: %f, loss: %f" % (
|
|
step, bleu_score, total_loss / self.save_step))
|
|
eval_summary = tf.Summary(value=[tf.Summary.Value(
|
|
tag='bleu', simple_value=bleu_score)])
|
|
self.log_writter.add_summary(eval_summary, step)
|
|
if step % self.save_step == 0:
|
|
total_loss = 0
|
|
|
|
|
|
def eval(self, train_step):
|
|
with self.eval_graph.as_default():
|
|
self.eval_saver.restore(self.eval_session, self.model_file)
|
|
bleu_score = 0
|
|
target_results = []
|
|
output_results = []
|
|
for step in range(0, self.eval_reader.data_size):
|
|
data = next(self.eval_data)
|
|
in_seq = data['in_seq']
|
|
in_seq_len = data['in_seq_len']
|
|
target_seq = data['target_seq']
|
|
target_seq_len = data['target_seq_len']
|
|
outputs = self.eval_session.run(
|
|
self.eval_output,
|
|
feed_dict={
|
|
self.eval_in_seq: in_seq,
|
|
self.eval_in_seq_len: in_seq_len})
|
|
for i in range(len(outputs)):
|
|
output = outputs[i]
|
|
target = target_seq[i]
|
|
output_text = reader.decode_text(output,
|
|
self.eval_reader.vocabs).split(' ')
|
|
target_text = reader.decode_text(target[1:],
|
|
self.eval_reader.vocabs).split(' ')
|
|
prob = int(self.eval_reader.data_size * self.batch_size / 10)
|
|
target_results.append([target_text])
|
|
output_results.append(output_text)
|
|
if random.randint(1, prob) == 1:
|
|
print('====================')
|
|
input_text = reader.decode_text(in_seq[i],
|
|
self.eval_reader.vocabs)
|
|
print('src:' + input_text)
|
|
print('output: ' + ' '.join(output_text))
|
|
print('target: ' + ' '.join(target_text))
|
|
return bleu.compute_bleu(target_results, output_results)[0] * 100
|
|
|
|
|
|
def reload_infer_model(self):
|
|
with self.infer_graph.as_default():
|
|
self.infer_saver.restore(self.infer_session, self.model_file)
|
|
|
|
|
|
def infer(self, text):
|
|
if not self.init_infer:
|
|
raise Exception('Infer graph is not inited!')
|
|
with self.infer_graph.as_default():
|
|
in_seq = reader.encode_text(text.split(' ') + ['</s>',],
|
|
self.infer_vocab_indices)
|
|
in_seq_len = len(in_seq)
|
|
outputs = self.infer_session.run(self.infer_output,
|
|
feed_dict={
|
|
self.infer_in_seq: [in_seq],
|
|
self.infer_in_seq_len: [in_seq_len]})
|
|
output = outputs[0]
|
|
output_text = reader.decode_text(output, self.infer_vocabs)
|
|
return output_text
|