# Copyright 2015 Google Inc. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Neural GPU.""" from __future__ import print_function import math import os import random import sys import threading import time import numpy as np from six.moves import xrange import tensorflow as tf import program_utils import data_utils as data import neural_gpu as ngpu import wmt_utils as wmt tf.app.flags.DEFINE_float("lr", 0.1, "Learning rate.") tf.app.flags.DEFINE_float("init_weight", 0.8, "Initial weights deviation.") tf.app.flags.DEFINE_float("max_grad_norm", 4.0, "Clip gradients to this norm.") tf.app.flags.DEFINE_float("cutoff", 1.2, "Cutoff at the gates.") tf.app.flags.DEFINE_float("curriculum_ppx", 9.9, "Move curriculum if ppl < X.") tf.app.flags.DEFINE_float("curriculum_seq", 0.3, "Move curriculum if seq < X.") tf.app.flags.DEFINE_float("dropout", 0.1, "Dropout that much.") tf.app.flags.DEFINE_float("grad_noise_scale", 0.0, "Gradient noise scale.") tf.app.flags.DEFINE_float("max_sampling_rate", 0.1, "Maximal sampling rate.") tf.app.flags.DEFINE_float("length_norm", 0.0, "Length normalization.") tf.app.flags.DEFINE_float("train_beam_freq", 0.0, "Beam-based training.") tf.app.flags.DEFINE_float("train_beam_anneal", 20000, "How many steps anneal.") tf.app.flags.DEFINE_integer("eval_beam_steps", 4, "How many beam steps eval.") tf.app.flags.DEFINE_integer("batch_size", 32, "Batch size.") tf.app.flags.DEFINE_integer("steps_per_checkpoint", 100, "Steps per epoch.") tf.app.flags.DEFINE_integer("nmaps", 64, "Number of floats in each cell.") tf.app.flags.DEFINE_integer("vec_size", 64, "Size of word vectors.") tf.app.flags.DEFINE_integer("train_data_size", 1000, "Training examples/len.") tf.app.flags.DEFINE_integer("max_length", 40, "Maximum length.") tf.app.flags.DEFINE_integer("random_seed", 125459, "Random seed.") tf.app.flags.DEFINE_integer("nconvs", 2, "How many convolutions / 1 step.") tf.app.flags.DEFINE_integer("kw", 3, "Kernel width.") tf.app.flags.DEFINE_integer("kh", 3, "Kernel height.") tf.app.flags.DEFINE_integer("height", 4, "Height.") tf.app.flags.DEFINE_integer("mem_size", -1, "Memory size (sqrt)") tf.app.flags.DEFINE_integer("soft_mem_size", 1024, "Softmax memory this size.") tf.app.flags.DEFINE_integer("num_gpus", 1, "Number of GPUs to use.") tf.app.flags.DEFINE_integer("num_replicas", 1, "Number of replicas in use.") tf.app.flags.DEFINE_integer("beam_size", 1, "Beam size during decoding. " "If 0, no decoder, the non-extended Neural GPU.") tf.app.flags.DEFINE_integer("max_target_vocab", 0, "Maximal size of target vocabulary.") tf.app.flags.DEFINE_integer("decode_offset", 0, "Offset for decoding.") tf.app.flags.DEFINE_integer("task", -1, "Task id when running on borg.") tf.app.flags.DEFINE_integer("nprint", 0, "How many test examples to print out.") tf.app.flags.DEFINE_integer("eval_bin_print", 3, "How many bins step in eval.") tf.app.flags.DEFINE_integer("mode", 0, "Mode: 0-train other-decode.") tf.app.flags.DEFINE_bool("atrous", False, "Whether to use atrous convs.") tf.app.flags.DEFINE_bool("layer_norm", False, "Do layer normalization.") tf.app.flags.DEFINE_bool("quantize", False, "Whether to quantize variables.") tf.app.flags.DEFINE_bool("do_train", True, "If false, only update memory.") tf.app.flags.DEFINE_bool("rnn_baseline", False, "If true build an RNN instead.") tf.app.flags.DEFINE_bool("simple_tokenizer", False, "If true, tokenize on spaces only, digits are 0.") tf.app.flags.DEFINE_bool("normalize_digits", True, "Whether to normalize digits with simple tokenizer.") tf.app.flags.DEFINE_integer("vocab_size", 16, "Joint vocabulary size.") tf.app.flags.DEFINE_string("data_dir", "/tmp", "Data directory") tf.app.flags.DEFINE_string("train_dir", "/tmp/", "Directory to store models.") tf.app.flags.DEFINE_string("test_file_prefix", "", "Files to test (.en,.fr).") tf.app.flags.DEFINE_integer("max_train_data_size", 0, "Limit on the size of training data (0: no limit).") tf.app.flags.DEFINE_string("word_vector_file_en", "", "Optional file with word vectors to start training.") tf.app.flags.DEFINE_string("word_vector_file_fr", "", "Optional file with word vectors to start training.") tf.app.flags.DEFINE_string("problem", "wmt", "What problem are we solving?.") tf.app.flags.DEFINE_integer("ps_tasks", 0, "Number of ps tasks used.") tf.app.flags.DEFINE_string("master", "", "Name of the TensorFlow master.") FLAGS = tf.app.flags.FLAGS EXTRA_EVAL = 10 EVAL_LEN_INCR = 8 MAXLEN_F = 2.0 def zero_split(tok_list, append=None): """Split tok_list (list of ints) on 0s, append int to all parts if given.""" res, cur, l = [], [], 0 for tok in tok_list: if tok == 0: if append is not None: cur.append(append) res.append(cur) l = max(l, len(cur)) cur = [] else: cur.append(tok) if append is not None: cur.append(append) res.append(cur) l = max(l, len(cur)) return res, l def read_data(source_path, target_path, buckets, max_size=None, print_out=True): """Read data from source and target files and put into buckets. Args: source_path: path to the files with token-ids for the source language. target_path: path to the file with token-ids for the target language; it must be aligned with the source file: n-th line contains the desired output for n-th line from the source_path. buckets: the buckets to use. max_size: maximum number of lines to read, all other will be ignored; if 0 or None, data files will be read completely (no limit). If set to 1, no data will be returned (empty lists of the right form). print_out: whether to print out status or not. Returns: data_set: a list of length len(_buckets); data_set[n] contains a list of (source, target) pairs read from the provided data files that fit into the n-th bucket, i.e., such that len(source) < _buckets[n][0] and len(target) < _buckets[n][1]; source and target are lists of token-ids. """ data_set = [[] for _ in buckets] counter = 0 if max_size != 1: with tf.gfile.GFile(source_path, mode="r") as source_file: with tf.gfile.GFile(target_path, mode="r") as target_file: source, target = source_file.readline(), target_file.readline() while source and target and (not max_size or counter < max_size): counter += 1 if counter % 100000 == 0 and print_out: print(" reading data line %d" % counter) sys.stdout.flush() source_ids = [int(x) for x in source.split()] target_ids = [int(x) for x in target.split()] source_ids, source_len = zero_split(source_ids) target_ids, target_len = zero_split(target_ids, append=wmt.EOS_ID) for bucket_id, size in enumerate(buckets): if source_len <= size and target_len <= size: data_set[bucket_id].append([source_ids, target_ids]) break source, target = source_file.readline(), target_file.readline() return data_set global_train_set = {"wmt": []} train_buckets_scale = {"wmt": []} def calculate_buckets_scale(data_set, buckets, problem): """Calculate buckets scales for the given data set.""" train_bucket_sizes = [len(data_set[b]) for b in xrange(len(buckets))] train_total_size = max(1, float(sum(train_bucket_sizes))) # A bucket scale is a list of increasing numbers from 0 to 1 that we'll use # to select a bucket. Length of [scale[i], scale[i+1]] is proportional to # the size if i-th training bucket, as used later. if problem not in train_buckets_scale: train_buckets_scale[problem] = [] train_buckets_scale[problem].append( [sum(train_bucket_sizes[:i + 1]) / train_total_size for i in xrange(len(train_bucket_sizes))]) return train_total_size def read_data_into_global(source_path, target_path, buckets, max_size=None, print_out=True): """Read data into the global variables (can be in a separate thread).""" # pylint: disable=global-variable-not-assigned global global_train_set, train_buckets_scale # pylint: enable=global-variable-not-assigned data_set = read_data(source_path, target_path, buckets, max_size, print_out) global_train_set["wmt"].append(data_set) train_total_size = calculate_buckets_scale(data_set, buckets, "wmt") if print_out: print(" Finished global data reading (%d)." % train_total_size) def initialize(sess=None): """Initialize data and model.""" global MAXLEN_F # Create training directory if it does not exist. if not tf.gfile.IsDirectory(FLAGS.train_dir): data.print_out("Creating training directory %s." % FLAGS.train_dir) tf.gfile.MkDir(FLAGS.train_dir) decode_suffix = "beam%dln%d" % (FLAGS.beam_size, int(100 * FLAGS.length_norm)) if FLAGS.mode == 0: decode_suffix = "" if FLAGS.task >= 0: data.log_filename = os.path.join(FLAGS.train_dir, "log%d%s" % (FLAGS.task, decode_suffix)) else: data.log_filename = os.path.join(FLAGS.train_dir, "neural_gpu/log") # Set random seed. if FLAGS.random_seed > 0: seed = FLAGS.random_seed + max(0, FLAGS.task) tf.set_random_seed(seed) random.seed(seed) np.random.seed(seed) # Check data sizes. assert data.bins max_length = min(FLAGS.max_length, data.bins[-1]) while len(data.bins) > 1 and data.bins[-2] >= max_length + EXTRA_EVAL: data.bins = data.bins[:-1] if sess is None and FLAGS.task == 0 and FLAGS.num_replicas > 1: if max_length > 60: max_length = max_length * 1 / 2 # Save memory on chief. min_length = min(14, max_length - 3) if FLAGS.problem == "wmt" else 3 for p in FLAGS.problem.split("-"): if p in ["progeval", "progsynth"]: min_length = max(26, min_length) assert max_length + 1 > min_length while len(data.bins) > 1 and data.bins[-2] >= max_length + EXTRA_EVAL: data.bins = data.bins[:-1] # Create checkpoint directory if it does not exist. if FLAGS.mode == 0 or FLAGS.task < 0: checkpoint_dir = os.path.join(FLAGS.train_dir, "neural_gpu%s" % ("" if FLAGS.task < 0 else str(FLAGS.task))) else: checkpoint_dir = FLAGS.train_dir if not tf.gfile.IsDirectory(checkpoint_dir): data.print_out("Creating checkpoint directory %s." % checkpoint_dir) tf.gfile.MkDir(checkpoint_dir) # Prepare data. if FLAGS.problem == "wmt": # Prepare WMT data. data.print_out("Preparing WMT data in %s" % FLAGS.data_dir) if FLAGS.simple_tokenizer: MAXLEN_F = 3.5 (en_train, fr_train, en_dev, fr_dev, en_path, fr_path) = wmt.prepare_wmt_data( FLAGS.data_dir, FLAGS.vocab_size, tokenizer=wmt.space_tokenizer, normalize_digits=FLAGS.normalize_digits) else: (en_train, fr_train, en_dev, fr_dev, en_path, fr_path) = wmt.prepare_wmt_data( FLAGS.data_dir, FLAGS.vocab_size) # Read data into buckets and compute their sizes. fr_vocab, rev_fr_vocab = wmt.initialize_vocabulary(fr_path) data.vocab = fr_vocab data.rev_vocab = rev_fr_vocab data.print_out("Reading development and training data (limit: %d)." % FLAGS.max_train_data_size) dev_set = {} dev_set["wmt"] = read_data(en_dev, fr_dev, data.bins) def data_read(size, print_out): read_data_into_global(en_train, fr_train, data.bins, size, print_out) data_read(50000, False) read_thread_small = threading.Thread( name="reading-data-small", target=lambda: data_read(900000, False)) read_thread_small.start() read_thread_full = threading.Thread( name="reading-data-full", target=lambda: data_read(FLAGS.max_train_data_size, True)) read_thread_full.start() data.print_out("Data reading set up.") else: # Prepare algorithmic data. en_path, fr_path = None, None tasks = FLAGS.problem.split("-") data_size = FLAGS.train_data_size for t in tasks: data.print_out("Generating data for %s." % t) if t in ["progeval", "progsynth"]: data.init_data(t, data.bins[-1], 20 * data_size, FLAGS.vocab_size) if len(program_utils.prog_vocab) > FLAGS.vocab_size - 2: raise ValueError("Increase vocab_size to %d for prog-tasks." % (len(program_utils.prog_vocab) + 2)) data.rev_vocab = program_utils.prog_vocab data.vocab = program_utils.prog_rev_vocab else: for l in xrange(max_length + EXTRA_EVAL - 1): data.init_data(t, l, data_size, FLAGS.vocab_size) data.init_data(t, data.bins[-2], data_size, FLAGS.vocab_size) data.init_data(t, data.bins[-1], data_size, FLAGS.vocab_size) if t not in global_train_set: global_train_set[t] = [] global_train_set[t].append(data.train_set[t]) calculate_buckets_scale(data.train_set[t], data.bins, t) dev_set = data.test_set # Grid-search parameters. lr = FLAGS.lr init_weight = FLAGS.init_weight max_grad_norm = FLAGS.max_grad_norm if sess is not None and FLAGS.task > -1: def job_id_factor(step): """If jobid / step mod 3 is 0, 1, 2: say 0, 1, -1.""" return ((((FLAGS.task / step) % 3) + 1) % 3) - 1 lr *= math.pow(2, job_id_factor(1)) init_weight *= math.pow(1.5, job_id_factor(3)) max_grad_norm *= math.pow(2, job_id_factor(9)) # Print out parameters. curriculum = FLAGS.curriculum_seq msg1 = ("layers %d kw %d h %d kh %d batch %d noise %.2f" % (FLAGS.nconvs, FLAGS.kw, FLAGS.height, FLAGS.kh, FLAGS.batch_size, FLAGS.grad_noise_scale)) msg2 = ("cut %.2f lr %.3f iw %.2f cr %.2f nm %d d%.4f gn %.2f %s" % (FLAGS.cutoff, lr, init_weight, curriculum, FLAGS.nmaps, FLAGS.dropout, max_grad_norm, msg1)) data.print_out(msg2) # Create model and initialize it. tf.get_variable_scope().set_initializer( tf.orthogonal_initializer(gain=1.8 * init_weight)) max_sampling_rate = FLAGS.max_sampling_rate if FLAGS.mode == 0 else 0.0 o = FLAGS.vocab_size if FLAGS.max_target_vocab < 1 else FLAGS.max_target_vocab ngpu.CHOOSE_K = FLAGS.soft_mem_size do_beam_model = FLAGS.train_beam_freq > 0.0001 and FLAGS.beam_size > 1 beam_size = FLAGS.beam_size if FLAGS.mode > 0 and not do_beam_model else 1 beam_size = min(beam_size, FLAGS.beam_size) beam_model = None def make_ngpu(cur_beam_size, back): return ngpu.NeuralGPU( FLAGS.nmaps, FLAGS.vec_size, FLAGS.vocab_size, o, FLAGS.dropout, max_grad_norm, FLAGS.cutoff, FLAGS.nconvs, FLAGS.kw, FLAGS.kh, FLAGS.height, FLAGS.mem_size, lr / math.sqrt(FLAGS.num_replicas), min_length + 3, FLAGS.num_gpus, FLAGS.num_replicas, FLAGS.grad_noise_scale, max_sampling_rate, atrous=FLAGS.atrous, do_rnn=FLAGS.rnn_baseline, do_layer_norm=FLAGS.layer_norm, beam_size=cur_beam_size, backward=back) if sess is None: with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)): model = make_ngpu(beam_size, True) if do_beam_model: tf.get_variable_scope().reuse_variables() beam_model = make_ngpu(FLAGS.beam_size, False) else: model = make_ngpu(beam_size, True) if do_beam_model: tf.get_variable_scope().reuse_variables() beam_model = make_ngpu(FLAGS.beam_size, False) sv = None if sess is None: # The supervisor configuration has a few overriden options. sv = tf.train.Supervisor(logdir=checkpoint_dir, is_chief=(FLAGS.task < 1), saver=model.saver, summary_op=None, save_summaries_secs=60, save_model_secs=15 * 60, global_step=model.global_step) config = tf.ConfigProto(allow_soft_placement=True) sess = sv.PrepareSession(FLAGS.master, config=config) data.print_out("Created model. Checkpoint dir %s" % checkpoint_dir) # Load model from parameters if a checkpoint exists. ckpt = tf.train.get_checkpoint_state(checkpoint_dir) if ckpt and tf.gfile.Exists(ckpt.model_checkpoint_path + ".index"): data.print_out("Reading model parameters from %s" % ckpt.model_checkpoint_path) model.saver.restore(sess, ckpt.model_checkpoint_path) elif sv is None: sess.run(tf.global_variables_initializer()) data.print_out("Initialized variables (no supervisor mode).") elif FLAGS.task < 1 and FLAGS.mem_size > 0: # sess.run(model.mem_norm_op) data.print_out("Created new model and normalized mem (on chief).") # Return the model and needed variables. return (model, beam_model, min_length, max_length, checkpoint_dir, (global_train_set, dev_set, en_path, fr_path), sv, sess) def m_step(model, beam_model, sess, batch_size, inp, target, bucket, nsteps, p): """Evaluation multi-step for program synthesis.""" state, scores, hist = None, [[-11.0 for _ in xrange(batch_size)]], [] for _ in xrange(nsteps): # Get the best beam (no training, just forward model). new_target, new_first, new_inp, new_scores = get_best_beam( beam_model, sess, inp, target, batch_size, FLAGS.beam_size, bucket, hist, p, test_mode=True) hist.append(new_first) _, _, _, state = model.step(sess, inp, new_target, False, state=state) inp = new_inp scores.append([max(scores[-1][i], new_scores[i]) for i in xrange(batch_size)]) # The final step with the true target. loss, res, _, _ = model.step(sess, inp, target, False, state=state) return loss, res, new_target, scores[1:] def single_test(bin_id, model, sess, nprint, batch_size, dev, p, print_out=True, offset=None, beam_model=None): """Test model on test data of length l using the given session.""" if not dev[p][bin_id]: data.print_out(" bin %d (%d)\t%s\tppl NA errors NA seq-errors NA" % (bin_id, data.bins[bin_id], p)) return 1.0, 1.0, 0.0 inpt, target = data.get_batch( bin_id, batch_size, dev[p], FLAGS.height, offset) if FLAGS.beam_size > 1 and beam_model: loss, res, new_tgt, scores = m_step( model, beam_model, sess, batch_size, inpt, target, bin_id, FLAGS.eval_beam_steps, p) score_avgs = [sum(s) / float(len(s)) for s in scores] score_maxs = [max(s) for s in scores] score_str = ["(%.2f, %.2f)" % (score_avgs[i], score_maxs[i]) for i in xrange(FLAGS.eval_beam_steps)] data.print_out(" == scores (avg, max): %s" % "; ".join(score_str)) errors, total, seq_err = data.accuracy(inpt, res, target, batch_size, nprint, new_tgt, scores[-1]) else: loss, res, _, _ = model.step(sess, inpt, target, False) errors, total, seq_err = data.accuracy(inpt, res, target, batch_size, nprint) seq_err = float(seq_err) / batch_size if total > 0: errors = float(errors) / total if print_out: data.print_out(" bin %d (%d)\t%s\tppl %.2f errors %.2f seq-errors %.2f" % (bin_id, data.bins[bin_id], p, data.safe_exp(loss), 100 * errors, 100 * seq_err)) return (errors, seq_err, loss) def assign_vectors(word_vector_file, embedding_key, vocab_path, sess): """Assign the embedding_key variable from the given word vectors file.""" # For words in the word vector file, set their embedding at start. if not tf.gfile.Exists(word_vector_file): data.print_out("Word vector file does not exist: %s" % word_vector_file) sys.exit(1) vocab, _ = wmt.initialize_vocabulary(vocab_path) vectors_variable = [v for v in tf.trainable_variables() if embedding_key == v.name] if len(vectors_variable) != 1: data.print_out("Word vector variable not found or too many.") sys.exit(1) vectors_variable = vectors_variable[0] vectors = vectors_variable.eval() data.print_out("Pre-setting word vectors from %s" % word_vector_file) with tf.gfile.GFile(word_vector_file, mode="r") as f: # Lines have format: dog 0.045123 -0.61323 0.413667 ... for line in f: line_parts = line.split() # The first part is the word. word = line_parts[0] if word in vocab: # Remaining parts are components of the vector. word_vector = np.array(map(float, line_parts[1:])) if len(word_vector) != FLAGS.vec_size: data.print_out("Warn: Word '%s', Expecting vector size %d, " "found %d" % (word, FLAGS.vec_size, len(word_vector))) else: vectors[vocab[word]] = word_vector # Assign the modified vectors to the vectors_variable in the graph. sess.run([vectors_variable.initializer], {vectors_variable.initializer.inputs[1]: vectors}) def print_vectors(embedding_key, vocab_path, word_vector_file): """Print vectors from the given variable.""" _, rev_vocab = wmt.initialize_vocabulary(vocab_path) vectors_variable = [v for v in tf.trainable_variables() if embedding_key == v.name] if len(vectors_variable) != 1: data.print_out("Word vector variable not found or too many.") sys.exit(1) vectors_variable = vectors_variable[0] vectors = vectors_variable.eval() l, s = vectors.shape[0], vectors.shape[1] data.print_out("Printing %d word vectors from %s to %s." % (l, embedding_key, word_vector_file)) with tf.gfile.GFile(word_vector_file, mode="w") as f: # Lines have format: dog 0.045123 -0.61323 0.413667 ... for i in xrange(l): f.write(rev_vocab[i]) for j in xrange(s): f.write(" %.8f" % vectors[i][j]) f.write("\n") def get_bucket_id(train_buckets_scale_c, max_cur_length, data_set): """Get a random bucket id.""" # Choose a bucket according to data distribution. Pick a random number # in [0, 1] and use the corresponding interval in train_buckets_scale. random_number_01 = np.random.random_sample() bucket_id = min([i for i in xrange(len(train_buckets_scale_c)) if train_buckets_scale_c[i] > random_number_01]) while bucket_id > 0 and not data_set[bucket_id]: bucket_id -= 1 for _ in xrange(10 if np.random.random_sample() < 0.9 else 1): if data.bins[bucket_id] > max_cur_length: random_number_01 = min(random_number_01, np.random.random_sample()) bucket_id = min([i for i in xrange(len(train_buckets_scale_c)) if train_buckets_scale_c[i] > random_number_01]) while bucket_id > 0 and not data_set[bucket_id]: bucket_id -= 1 return bucket_id def score_beams(beams, target, inp, history, p, print_out=False, test_mode=False): """Score beams.""" if p == "progsynth": return score_beams_prog(beams, target, inp, history, print_out, test_mode) elif test_mode: return beams[0], 10.0 if str(beams[0][:len(target)]) == str(target) else 0.0 else: history_s = [str(h) for h in history] best, best_score, tgt, eos_id = None, -1000.0, target, None if p == "wmt": eos_id = wmt.EOS_ID if eos_id and eos_id in target: tgt = target[:target.index(eos_id)] for beam in beams: if eos_id and eos_id in beam: beam = beam[:beam.index(eos_id)] l = min(len(tgt), len(beam)) score = len([i for i in xrange(l) if tgt[i] == beam[i]]) / float(len(tgt)) hist_score = 20.0 if str([b for b in beam if b > 0]) in history_s else 0.0 if score < 1.0: score -= hist_score if score > best_score: best = beam best_score = score return best, best_score def score_beams_prog(beams, target, inp, history, print_out=False, test_mode=False): """Score beams for program synthesis.""" tgt_prog = linearize(target, program_utils.prog_vocab, True, 1) hist_progs = [linearize(h, program_utils.prog_vocab, True, 1) for h in history] tgt_set = set(target) if print_out: print("target: ", tgt_prog) inps, tgt_outs = [], [] for i in xrange(3): ilist = [inp[i + 1, l] for l in xrange(inp.shape[1])] clist = [program_utils.prog_vocab[x] for x in ilist if x > 0] olist = clist[clist.index("]") + 1:] # outputs clist = clist[1:clist.index("]")] # inputs inps.append([int(x) for x in clist]) if olist[0] == "[": # olist may be [int] or just int tgt_outs.append(str([int(x) for x in olist[1:-1]])) else: if len(olist) == 1: tgt_outs.append(olist[0]) else: print([program_utils.prog_vocab[x] for x in ilist if x > 0]) print(olist) print(tgt_prog) print(program_utils.evaluate(tgt_prog, {"a": inps[-1]})) print("AAAAA") tgt_outs.append(olist[0]) if not test_mode: for _ in xrange(7): ilen = np.random.randint(len(target) - 3) + 1 inps.append([random.choice(range(-15, 15)) for _ in range(ilen)]) tgt_outs.extend([program_utils.evaluate(tgt_prog, {"a": inp}) for inp in inps[3:]]) best, best_prog, best_score = None, "", -1000.0 for beam in beams: b_prog = linearize(beam, program_utils.prog_vocab, True, 1) b_set = set(beam) jsim = len(tgt_set & b_set) / float(len(tgt_set | b_set)) b_outs = [program_utils.evaluate(b_prog, {"a": inp}) for inp in inps] errs = len([x for x in b_outs if x == "ERROR"]) imatches = len([i for i in xrange(3) if b_outs[i] == tgt_outs[i]]) perfect = 10.0 if imatches == 3 else 0.0 hist_score = 20.0 if b_prog in hist_progs else 0.0 if test_mode: score = perfect - errs else: matches = len([i for i in xrange(10) if b_outs[i] == tgt_outs[i]]) score = perfect + matches + jsim - errs if score < 10.0: score -= hist_score # print b_prog # print "jsim: ", jsim, " errs: ", errs, " mtchs: ", matches, " s: ", score if score > best_score: best = beam best_prog = b_prog best_score = score if print_out: print("best score: ", best_score, " best prog: ", best_prog) return best, best_score def get_best_beam(beam_model, sess, inp, target, batch_size, beam_size, bucket, history, p, test_mode=False): """Run beam_model, score beams, and return the best as target and in input.""" _, output_logits, _, _ = beam_model.step( sess, inp, target, None, beam_size=FLAGS.beam_size) new_targets, new_firsts, scores, new_inp = [], [], [], np.copy(inp) for b in xrange(batch_size): outputs = [] history_b = [[h[b, 0, l] for l in xrange(data.bins[bucket])] for h in history] for beam_idx in xrange(beam_size): outputs.append([int(o[beam_idx * batch_size + b]) for o in output_logits]) target_t = [target[b, 0, l] for l in xrange(data.bins[bucket])] best, best_score = score_beams( outputs, [t for t in target_t if t > 0], inp[b, :, :], [[t for t in h if t > 0] for h in history_b], p, test_mode=test_mode) scores.append(best_score) if 1 in best: # Only until _EOS. best = best[:best.index(1) + 1] best += [0 for _ in xrange(len(target_t) - len(best))] new_targets.append([best]) first, _ = score_beams( outputs, [t for t in target_t if t > 0], inp[b, :, :], [[t for t in h if t > 0] for h in history_b], p, test_mode=True) if 1 in first: # Only until _EOS. first = first[:first.index(1) + 1] first += [0 for _ in xrange(len(target_t) - len(first))] new_inp[b, 0, :] = np.array(first, dtype=np.int32) new_firsts.append([first]) # Change target if we found a great answer. new_target = np.array(new_targets, dtype=np.int32) for b in xrange(batch_size): if scores[b] >= 10.0: target[b, 0, :] = new_target[b, 0, :] new_first = np.array(new_firsts, dtype=np.int32) return new_target, new_first, new_inp, scores def train(): """Train the model.""" batch_size = FLAGS.batch_size * FLAGS.num_gpus (model, beam_model, min_length, max_length, checkpoint_dir, (train_set, dev_set, en_vocab_path, fr_vocab_path), sv, sess) = initialize() with sess.as_default(): quant_op = model.quantize_op max_cur_length = min(min_length + 3, max_length) prev_acc_perp = [1000000 for _ in xrange(5)] prev_seq_err = 1.0 is_chief = FLAGS.task < 1 do_report = False # Main traning loop. while not sv.ShouldStop(): global_step, max_cur_length, learning_rate = sess.run( [model.global_step, model.cur_length, model.lr]) acc_loss, acc_l1, acc_total, acc_errors, acc_seq_err = 0.0, 0.0, 0, 0, 0 acc_grad_norm, step_count, step_c1, step_time = 0.0, 0, 0, 0.0 # For words in the word vector file, set their embedding at start. bound1 = FLAGS.steps_per_checkpoint - 1 if FLAGS.word_vector_file_en and global_step < bound1 and is_chief: assign_vectors(FLAGS.word_vector_file_en, "embedding:0", en_vocab_path, sess) if FLAGS.max_target_vocab < 1: assign_vectors(FLAGS.word_vector_file_en, "target_embedding:0", en_vocab_path, sess) if FLAGS.word_vector_file_fr and global_step < bound1 and is_chief: assign_vectors(FLAGS.word_vector_file_fr, "embedding:0", fr_vocab_path, sess) if FLAGS.max_target_vocab < 1: assign_vectors(FLAGS.word_vector_file_fr, "target_embedding:0", fr_vocab_path, sess) for _ in xrange(FLAGS.steps_per_checkpoint): step_count += 1 step_c1 += 1 global_step = int(model.global_step.eval()) train_beam_anneal = global_step / float(FLAGS.train_beam_anneal) train_beam_freq = FLAGS.train_beam_freq * min(1.0, train_beam_anneal) p = random.choice(FLAGS.problem.split("-")) train_set = global_train_set[p][-1] bucket_id = get_bucket_id(train_buckets_scale[p][-1], max_cur_length, train_set) # Prefer longer stuff 60% of time if not wmt. if np.random.randint(100) < 60 and FLAGS.problem != "wmt": bucket1 = get_bucket_id(train_buckets_scale[p][-1], max_cur_length, train_set) bucket_id = max(bucket1, bucket_id) # Run a step and time it. start_time = time.time() inp, target = data.get_batch(bucket_id, batch_size, train_set, FLAGS.height) noise_param = math.sqrt(math.pow(global_step + 1, -0.55) * prev_seq_err) * FLAGS.grad_noise_scale # In multi-step mode, we use best from beam for middle steps. state, new_target, scores, history = None, None, None, [] while (FLAGS.beam_size > 1 and train_beam_freq > np.random.random_sample()): # Get the best beam (no training, just forward model). new_target, new_first, new_inp, scores = get_best_beam( beam_model, sess, inp, target, batch_size, FLAGS.beam_size, bucket_id, history, p) history.append(new_first) # Training step with the previous input and the best beam as target. _, _, _, state = model.step(sess, inp, new_target, FLAGS.do_train, noise_param, update_mem=True, state=state) # Change input to the new one for the next step. inp = new_inp # If all results are great, stop (todo: not to wait for all?). if FLAGS.nprint > 1: print(scores) if sum(scores) / float(len(scores)) >= 10.0: break # The final step with the true target. loss, res, gnorm, _ = model.step( sess, inp, target, FLAGS.do_train, noise_param, update_mem=True, state=state) step_time += time.time() - start_time acc_grad_norm += 0.0 if gnorm is None else float(gnorm) # Accumulate statistics. acc_loss += loss acc_l1 += loss errors, total, seq_err = data.accuracy( inp, res, target, batch_size, 0, new_target, scores) if FLAGS.nprint > 1: print("seq_err: ", seq_err) acc_total += total acc_errors += errors acc_seq_err += seq_err # Report summary every 10 steps. if step_count + 3 > FLAGS.steps_per_checkpoint: do_report = True # Don't polute plot too early. if is_chief and step_count % 10 == 1 and do_report: cur_loss = acc_l1 / float(step_c1) acc_l1, step_c1 = 0.0, 0 cur_perp = data.safe_exp(cur_loss) summary = tf.Summary() summary.value.extend( [tf.Summary.Value(tag="log_perplexity", simple_value=cur_loss), tf.Summary.Value(tag="perplexity", simple_value=cur_perp)]) sv.SummaryComputed(sess, summary, global_step) # Normalize and print out accumulated statistics. acc_loss /= step_count step_time /= FLAGS.steps_per_checkpoint acc_seq_err = float(acc_seq_err) / (step_count * batch_size) prev_seq_err = max(0.0, acc_seq_err - 0.02) # No noise at error < 2%. acc_errors = float(acc_errors) / acc_total if acc_total > 0 else 1.0 t_size = float(sum([len(x) for x in train_set])) / float(1000000) msg = ("step %d step-time %.2f train-size %.3f lr %.6f grad-norm %.4f" % (global_step + 1, step_time, t_size, learning_rate, acc_grad_norm / FLAGS.steps_per_checkpoint)) data.print_out("%s len %d ppl %.6f errors %.2f sequence-errors %.2f" % (msg, max_cur_length, data.safe_exp(acc_loss), 100*acc_errors, 100*acc_seq_err)) # If errors are below the curriculum threshold, move curriculum forward. is_good = FLAGS.curriculum_ppx > data.safe_exp(acc_loss) is_good = is_good and FLAGS.curriculum_seq > acc_seq_err if is_good and is_chief: if FLAGS.quantize: # Quantize weights. data.print_out(" Quantizing parameters.") sess.run([quant_op]) # Increase current length (until the next with training data). sess.run(model.cur_length_incr_op) # Forget last perplexities if we're not yet at the end. if max_cur_length < max_length: prev_acc_perp.append(1000000) # Lower learning rate if we're worse than the last 5 checkpoints. acc_perp = data.safe_exp(acc_loss) if acc_perp > max(prev_acc_perp[-5:]) and is_chief: sess.run(model.lr_decay_op) prev_acc_perp.append(acc_perp) # Save checkpoint. if is_chief: checkpoint_path = os.path.join(checkpoint_dir, "neural_gpu.ckpt") model.saver.save(sess, checkpoint_path, global_step=model.global_step) # Run evaluation. bin_bound = 4 for p in FLAGS.problem.split("-"): total_loss, total_err, tl_counter = 0.0, 0.0, 0 for bin_id in xrange(len(data.bins)): if bin_id < bin_bound or bin_id % FLAGS.eval_bin_print == 1: err, _, loss = single_test(bin_id, model, sess, FLAGS.nprint, batch_size * 4, dev_set, p, beam_model=beam_model) if loss > 0.0: total_loss += loss total_err += err tl_counter += 1 test_loss = total_loss / max(1, tl_counter) test_err = total_err / max(1, tl_counter) test_perp = data.safe_exp(test_loss) summary = tf.Summary() summary.value.extend( [tf.Summary.Value(tag="test/%s/loss" % p, simple_value=test_loss), tf.Summary.Value(tag="test/%s/error" % p, simple_value=test_err), tf.Summary.Value(tag="test/%s/perplexity" % p, simple_value=test_perp)]) sv.SummaryComputed(sess, summary, global_step) def linearize(output, rev_fr_vocab, simple_tokenizer=None, eos_id=wmt.EOS_ID): # If there is an EOS symbol in outputs, cut them at that point (WMT). if eos_id in output: output = output[:output.index(eos_id)] # Print out French sentence corresponding to outputs. if simple_tokenizer or FLAGS.simple_tokenizer: vlen = len(rev_fr_vocab) def vget(o): if o < vlen: return rev_fr_vocab[o] return "UNK" return " ".join([vget(o) for o in output]) else: return wmt.basic_detokenizer([rev_fr_vocab[o] for o in output]) def evaluate(): """Evaluate an existing model.""" batch_size = FLAGS.batch_size * FLAGS.num_gpus with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess: (model, beam_model, _, _, _, (_, dev_set, en_vocab_path, fr_vocab_path), _, sess) = initialize(sess) for p in FLAGS.problem.split("-"): for bin_id in xrange(len(data.bins)): if (FLAGS.task >= 0 and bin_id > 4) or (FLAGS.nprint == 0 and bin_id > 8 and p == "wmt"): break single_test(bin_id, model, sess, FLAGS.nprint, batch_size, dev_set, p, beam_model=beam_model) path = FLAGS.test_file_prefix xid = "" if FLAGS.task < 0 else ("%.4d" % (FLAGS.task+FLAGS.decode_offset)) en_path, fr_path = path + ".en" + xid, path + ".fr" + xid # Evaluate the test file if they exist. if path and tf.gfile.Exists(en_path) and tf.gfile.Exists(fr_path): data.print_out("Translating test set %s" % en_path) # Read lines. en_lines, fr_lines = [], [] with tf.gfile.GFile(en_path, mode="r") as f: for line in f: en_lines.append(line.strip()) with tf.gfile.GFile(fr_path, mode="r") as f: for line in f: fr_lines.append(line.strip()) # Tokenize and convert to ids. en_vocab, _ = wmt.initialize_vocabulary(en_vocab_path) _, rev_fr_vocab = wmt.initialize_vocabulary(fr_vocab_path) if FLAGS.simple_tokenizer: en_ids = [wmt.sentence_to_token_ids( l, en_vocab, tokenizer=wmt.space_tokenizer, normalize_digits=FLAGS.normalize_digits) for l in en_lines] else: en_ids = [wmt.sentence_to_token_ids(l, en_vocab) for l in en_lines] # Translate. results = [] for idx, token_ids in enumerate(en_ids): if idx % 5 == 0: data.print_out("Translating example %d of %d." % (idx, len(en_ids))) # Which bucket does it belong to? buckets = [b for b in xrange(len(data.bins)) if data.bins[b] >= len(token_ids)] if buckets: result, result_cost = [], 100000000.0 for bucket_id in buckets: if data.bins[bucket_id] > MAXLEN_F * len(token_ids) + EVAL_LEN_INCR: break # Get a 1-element batch to feed the sentence to the model. used_batch_size = 1 # batch_size inp, target = data.get_batch( bucket_id, used_batch_size, None, FLAGS.height, preset=([token_ids], [[]])) loss, output_logits, _, _ = model.step( sess, inp, target, None, beam_size=FLAGS.beam_size) outputs = [int(o[0]) for o in output_logits] loss = loss[0] - (data.bins[bucket_id] * FLAGS.length_norm) if FLAGS.simple_tokenizer: cur_out = outputs if wmt.EOS_ID in cur_out: cur_out = cur_out[:cur_out.index(wmt.EOS_ID)] res_tags = [rev_fr_vocab[o] for o in cur_out] bad_words, bad_brack = wmt.parse_constraints(token_ids, res_tags) loss += 1000.0 * bad_words + 100.0 * bad_brack # print (bucket_id, loss) if loss < result_cost: result = outputs result_cost = loss final = linearize(result, rev_fr_vocab) results.append("%s\t%s\n" % (final, fr_lines[idx])) # print result_cost sys.stderr.write(results[-1]) sys.stderr.flush() else: sys.stderr.write("TOOO_LONG\t%s\n" % fr_lines[idx]) sys.stderr.flush() if xid: decode_suffix = "beam%dln%dn" % (FLAGS.beam_size, int(100 * FLAGS.length_norm)) with tf.gfile.GFile(path + ".res" + decode_suffix + xid, mode="w") as f: for line in results: f.write(line) def mul(l): res = 1.0 for s in l: res *= s return res def interactive(): """Interactively probe an existing model.""" with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess: # Initialize model. (model, _, _, _, _, (_, _, en_path, fr_path), _, _) = initialize(sess) # Load vocabularies. en_vocab, rev_en_vocab = wmt.initialize_vocabulary(en_path) _, rev_fr_vocab = wmt.initialize_vocabulary(fr_path) # Print out vectors and variables. if FLAGS.nprint > 0 and FLAGS.word_vector_file_en: print_vectors("embedding:0", en_path, FLAGS.word_vector_file_en) if FLAGS.nprint > 0 and FLAGS.word_vector_file_fr: print_vectors("target_embedding:0", fr_path, FLAGS.word_vector_file_fr) total = 0 for v in tf.trainable_variables(): shape = v.get_shape().as_list() total += mul(shape) print(v.name, shape, mul(shape)) print(total) # Start interactive loop. sys.stdout.write("Input to Neural GPU Translation Model.\n") sys.stdout.write("> ") sys.stdout.flush() inpt = sys.stdin.readline(), "" while inpt: cures = [] # Get token-ids for the input sentence. if FLAGS.simple_tokenizer: token_ids = wmt.sentence_to_token_ids( inpt, en_vocab, tokenizer=wmt.space_tokenizer, normalize_digits=FLAGS.normalize_digits) else: token_ids = wmt.sentence_to_token_ids(inpt, en_vocab) print([rev_en_vocab[t] for t in token_ids]) # Which bucket does it belong to? buckets = [b for b in xrange(len(data.bins)) if data.bins[b] >= max(len(token_ids), len(cures))] if cures: buckets = [buckets[0]] if buckets: result, result_cost = [], 10000000.0 for bucket_id in buckets: if data.bins[bucket_id] > MAXLEN_F * len(token_ids) + EVAL_LEN_INCR: break glen = 1 for gen_idx in xrange(glen): # Get a 1-element batch to feed the sentence to the model. inp, target = data.get_batch( bucket_id, 1, None, FLAGS.height, preset=([token_ids], [cures])) loss, output_logits, _, _ = model.step( sess, inp, target, None, beam_size=FLAGS.beam_size, update_mem=False) # If it is a greedy decoder, outputs are argmaxes of output_logits. if FLAGS.beam_size > 1: outputs = [int(o) for o in output_logits] else: loss = loss[0] - (data.bins[bucket_id] * FLAGS.length_norm) outputs = [int(np.argmax(logit, axis=1)) for logit in output_logits] print([rev_fr_vocab[t] for t in outputs]) print(loss, data.bins[bucket_id]) print(linearize(outputs, rev_fr_vocab)) cures.append(outputs[gen_idx]) print(cures) print(linearize(cures, rev_fr_vocab)) if FLAGS.simple_tokenizer: cur_out = outputs if wmt.EOS_ID in cur_out: cur_out = cur_out[:cur_out.index(wmt.EOS_ID)] res_tags = [rev_fr_vocab[o] for o in cur_out] bad_words, bad_brack = wmt.parse_constraints(token_ids, res_tags) loss += 1000.0 * bad_words + 100.0 * bad_brack if loss < result_cost: result = outputs result_cost = loss print("FINAL", result_cost) print([rev_fr_vocab[t] for t in result]) print(linearize(result, rev_fr_vocab)) else: print("TOOO_LONG") sys.stdout.write("> ") sys.stdout.flush() inpt = sys.stdin.readline(), "" def main(_): if FLAGS.mode == 0: train() elif FLAGS.mode == 1: evaluate() else: interactive() if __name__ == "__main__": tf.app.run()