#!/usr/bin/env python3 # # Copyright 2016 Google Inc. All Rights Reserved. # Copyright 2017 Sourced Technologies S. L. # # 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. """Submatrix-wise Vector Embedding Learner. Implementation of SwiVel algorithm described at: http://arxiv.org/abs/1602.02215 This program expects an input directory that contains the following files. row_vocab.txt, col_vocab.txt The row an column vocabulary files. Each file should contain one token per line; these will be used to generate a tab-separate file containing the trained embeddings. row_sums.txt, col_sum.txt The matrix row and column marginal sums. Each file should contain one decimal floating point number per line which corresponds to the marginal count of the matrix for that row or column. shards.recs A file containing the sub-matrix shards, stored as TFRecords. Each shard is expected to be a serialzed tf.Example protocol buffer with the following properties: global_row: the global row indicies contained in the shard global_col: the global column indicies contained in the shard sparse_local_row, sparse_local_col, sparse_value: three parallel arrays that are a sparse representation of the submatrix counts. It will generate embeddings, training from the input directory for the specified number of epochs. When complete, it will output the trained vectors to a tab-separated file that contains one line per embedding. Row and column embeddings are stored in separate files. """ import glob import math import os import threading import time import numpy as np import tensorflow as tf from tensorflow.python.client import device_lib flags = tf.app.flags flags.DEFINE_string("input_base_path", None, "Directory containing input shards, vocabularies, " "and marginals.") flags.DEFINE_string("output_base_path", None, "Path where to write the trained embeddings.") flags.DEFINE_integer("embedding_size", 300, "Size of the embeddings") flags.DEFINE_boolean("trainable_bias", False, "Biases are trainable") flags.DEFINE_integer("submatrix_rows", 4096, "Rows in each training submatrix. This must match " "the training data.") flags.DEFINE_integer("submatrix_cols", 4096, "Rows in each training submatrix. This must match " "the training data.") flags.DEFINE_float("loss_multiplier", 1.0 / 4096, "constant multiplier on loss.") flags.DEFINE_float("confidence_exponent", 0.5, "Exponent for l2 confidence function") flags.DEFINE_float("confidence_scale", 0.25, "Scale for l2 confidence function") flags.DEFINE_float("confidence_base", 0.1, "Base for l2 confidence function") flags.DEFINE_float("learning_rate", 1.0, "Initial learning rate") flags.DEFINE_string("optimizer", "Adagrad", "SGD optimizer (tf.train.*Optimizer)") flags.DEFINE_integer("num_concurrent_steps", 2, "Number of threads to train with") flags.DEFINE_integer("num_readers", 4, "Number of threads to read the input data and feed it") flags.DEFINE_float("num_epochs", 40, "Number epochs to train for") flags.DEFINE_float("per_process_gpu_memory_fraction", 0, "Fraction of GPU memory to use, 0 means allow_growth") flags.DEFINE_integer("num_gpus", 0, "Number of GPUs to use, 0 means all available") flags.DEFINE_string("logs", "", "Path for TensorBoard logs (empty value disables them)") FLAGS = flags.FLAGS def log(message, *args, **kwargs): tf.logging.info(message, *args, **kwargs) def get_available_gpus(): return [d.name for d in device_lib.list_local_devices() if d.device_type == "GPU"] def embeddings_with_init(vocab_size, embedding_dim, name): """Creates and initializes the embedding tensors.""" return tf.get_variable(name=name, shape=[vocab_size, embedding_dim], initializer=tf.random_normal_initializer( stddev=math.sqrt(1.0 / embedding_dim))) def count_matrix_input(filenames, submatrix_rows, submatrix_cols): """Reads submatrix shards from disk.""" filename_queue = tf.train.string_input_producer(filenames) reader = tf.WholeFileReader() _, serialized_example = reader.read(filename_queue) features = tf.parse_single_example( serialized_example, features={ "global_row": tf.FixedLenFeature([submatrix_rows], dtype=tf.int64), "global_col": tf.FixedLenFeature([submatrix_cols], dtype=tf.int64), "sparse_local_row": tf.VarLenFeature(dtype=tf.int64), "sparse_local_col": tf.VarLenFeature(dtype=tf.int64), "sparse_value": tf.VarLenFeature(dtype=tf.float32) }) global_row = features["global_row"] global_col = features["global_col"] sparse_local_row = features["sparse_local_row"].values sparse_local_col = features["sparse_local_col"].values sparse_count = features["sparse_value"].values sparse_indices = tf.concat(axis=1, values=[tf.expand_dims(sparse_local_row, 1), tf.expand_dims(sparse_local_col, 1)]) count = tf.sparse_to_dense(sparse_indices, [submatrix_rows, submatrix_cols], sparse_count, validate_indices=False) queued_global_row, queued_global_col, queued_count = tf.train.batch( [global_row, global_col, count], batch_size=1, num_threads=FLAGS.num_readers, capacity=32) queued_global_row = tf.reshape(queued_global_row, [submatrix_rows]) queued_global_col = tf.reshape(queued_global_col, [submatrix_cols]) queued_count = tf.reshape(queued_count, [submatrix_rows, submatrix_cols]) return queued_global_row, queued_global_col, queued_count def read_marginals_file(filename): """Reads text file with one number per line to an array.""" with open(filename) as lines: return [float(line) for line in lines] def write_embedding_tensor_to_disk(vocab_path, output_path, sess, embedding): """Writes tensor to output_path as tsv""" # Fetch the embedding values from the model embeddings = sess.run(embedding) with open(output_path, "w") as out_f: with open(vocab_path) as vocab_f: for index, word in enumerate(vocab_f): word = word.strip() embedding = embeddings[index] out_f.write(word + "\t" + "\t".join( [str(x) for x in embedding]) + "\n") def write_embeddings_to_disk(config, model, sess): """Writes row and column embeddings disk""" # Row Embedding row_vocab_path = config.input_base_path + "/row_vocab.txt" row_embedding_output_path = config.output_base_path + "/row_embedding.tsv" log("Writing row embeddings to: %s", row_embedding_output_path) write_embedding_tensor_to_disk(row_vocab_path, row_embedding_output_path, sess, model.row_embedding) # Column Embedding col_vocab_path = config.input_base_path + "/col_vocab.txt" col_embedding_output_path = config.output_base_path + "/col_embedding.tsv" log("Writing column embeddings to: %s", col_embedding_output_path) write_embedding_tensor_to_disk(col_vocab_path, col_embedding_output_path, sess, model.col_embedding) class SwivelModel: """Small class to gather needed pieces from a Graph being built.""" def __init__(self, config): """Construct graph for dmc.""" self._config = config # Create paths to input data files log("Reading model from: %s", config.input_base_path) count_matrix_files = glob.glob(os.path.join(config.input_base_path, "shard-*.pb")) row_sums_path = os.path.join(config.input_base_path, "row_sums.txt") col_sums_path = os.path.join(config.input_base_path, "col_sums.txt") # Read marginals row_sums = read_marginals_file(row_sums_path) col_sums = read_marginals_file(col_sums_path) self.n_rows = len(row_sums) self.n_cols = len(col_sums) log("Matrix dim: (%d,%d) SubMatrix dim: (%d,%d)", self.n_rows, self.n_cols, config.submatrix_rows, config.submatrix_cols) if self.n_cols < config.submatrix_cols: raise ValueError( "submatrix_cols={0} can not be bigger than columns number={1} " "(specify submatrix_cols={1})".format(config.submatrix_cols, self.n_cols)) if self.n_rows < config.submatrix_rows: raise ValueError( "submatrix_rows={0} can not be bigger than rows number={1} " "(specify submatrix_rows={1})".format(config.submatrix_rows, self.n_cols)) self.n_submatrices = ( self.n_rows * self.n_cols / (config.submatrix_rows * config.submatrix_cols)) log("n_submatrices: %d", self.n_submatrices) with tf.device("/cpu:0"): # ===== CREATE VARIABLES ====== # Get input global_row, global_col, count = count_matrix_input( count_matrix_files, config.submatrix_rows, config.submatrix_cols) # Embeddings self.row_embedding = embeddings_with_init( embedding_dim=config.embedding_size, vocab_size=self.n_rows, name="row_embedding") self.col_embedding = embeddings_with_init( embedding_dim=config.embedding_size, vocab_size=self.n_cols, name="col_embedding") tf.summary.histogram("row_emb", self.row_embedding) tf.summary.histogram("col_emb", self.col_embedding) matrix_log_sum = math.log(np.sum(row_sums) + 1) row_bias_init = [math.log(x + 1) for x in row_sums] col_bias_init = [math.log(x + 1) for x in col_sums] self.row_bias = tf.Variable( row_bias_init, trainable=config.trainable_bias) self.col_bias = tf.Variable( col_bias_init, trainable=config.trainable_bias) tf.summary.histogram("row_bias", self.row_bias) tf.summary.histogram("col_bias", self.col_bias) # Add optimizer l2_losses = [] sigmoid_losses = [] self.global_step = tf.Variable(0, name="global_step") learning_rate = tf.Variable(config.learning_rate, name="learning_rate") opt = getattr(tf.train, FLAGS.optimizer + "Optimizer")( learning_rate) tf.summary.scalar("learning_rate", learning_rate) all_grads = [] devices = ["/gpu:%d" % i for i in range(FLAGS.num_gpus)] \ if FLAGS.num_gpus > 0 else get_available_gpus() self.devices_number = len(devices) if not self.devices_number: devices = ["/cpu:0"] self.devices_number = 1 for dev in devices: with tf.device(dev): with tf.name_scope(dev[1:].replace(":", "_")): # ===== CREATE GRAPH ===== # Fetch embeddings. selected_row_embedding = tf.nn.embedding_lookup( self.row_embedding, global_row) selected_col_embedding = tf.nn.embedding_lookup( self.col_embedding, global_col) # Fetch biases. selected_row_bias = tf.nn.embedding_lookup( [self.row_bias], global_row) selected_col_bias = tf.nn.embedding_lookup( [self.col_bias], global_col) # Multiply the row and column embeddings to generate # predictions. predictions = tf.matmul( selected_row_embedding, selected_col_embedding, transpose_b=True) # These binary masks separate zero from non-zero values. count_is_nonzero = tf.to_float(tf.cast(count, tf.bool)) count_is_zero = 1 - count_is_nonzero objectives = count_is_nonzero * tf.log(count + 1e-30) objectives -= tf.reshape( selected_row_bias, [config.submatrix_rows, 1]) objectives -= selected_col_bias objectives += matrix_log_sum err = predictions - objectives # The confidence function scales the L2 loss based on # the raw co-occurrence count. l2_confidence = config.confidence_base + config.confidence_scale * tf.pow( count, config.confidence_exponent) l2_loss = config.loss_multiplier * tf.reduce_sum( 0.5 * l2_confidence * err * err * count_is_nonzero) l2_losses.append(tf.expand_dims(l2_loss, 0)) sigmoid_loss = config.loss_multiplier * tf.reduce_sum( tf.nn.softplus(err) * count_is_zero) sigmoid_losses.append(tf.expand_dims(sigmoid_loss, 0)) loss = l2_loss + sigmoid_loss grads = opt.compute_gradients(loss) all_grads.append(grads) with tf.device("/cpu:0"): # ===== MERGE LOSSES ===== l2_loss = tf.reduce_mean(tf.concat(axis=0, values=l2_losses), 0, name="l2_loss") sigmoid_loss = tf.reduce_mean( tf.concat(axis=0, values=sigmoid_losses), 0, name="sigmoid_loss") overall_loss = l2_loss + sigmoid_loss average = tf.train.ExponentialMovingAverage(0.999) loss_average_op = average.apply( (overall_loss, l2_loss, sigmoid_loss)) self.loss = average.average(overall_loss) tf.summary.scalar("overall_loss", self.loss) tf.summary.scalar("l2_loss", average.average(l2_loss)) tf.summary.scalar("sigmoid_loss", average.average(sigmoid_loss)) # Apply the gradients to adjust the shared variables. apply_gradient_ops = [] for grads in all_grads: apply_gradient_ops.append(opt.apply_gradients( grads, global_step=self.global_step)) self.train_op = tf.group(loss_average_op, *apply_gradient_ops) self.saver = tf.train.Saver(sharded=True) def initialize_summary(self, sess): log("creating TensorBoard stuff...") self.summary = tf.summary.merge_all() self.writer = tf.summary.FileWriter(FLAGS.logs, sess.graph) projector_config = \ tf.contrib.tensorboard.plugins.projector.ProjectorConfig() embedding_config = projector_config.embeddings.add() length = min(10000, self.n_rows, self.n_cols) self.embedding10k = tf.Variable( tf.zeros((length, self._config.embedding_size)), name="top10k_embedding") embedding_config.tensor_name = self.embedding10k.name embedding_config.metadata_path = os.path.join( self._config.input_base_path, "row_vocab.txt") tf.contrib.tensorboard.plugins.projector.visualize_embeddings( self.writer, projector_config) self.saver = tf.train.Saver((self.embedding10k,), max_to_keep=1) def write_summary(self, sess): log("writing the summary...") length = min(10000, self.n_rows, self.n_cols) assignment = self.embedding10k.assign( (self.row_embedding[:length] + self.col_embedding[:length]) / 2) summary, _, global_step = sess.run( (self.summary, assignment, self.global_step)) self.writer.add_summary(summary, global_step) self.saver.save( sess, os.path.join(FLAGS.logs, "embeddings10k.checkpoint"), global_step) def main(_): tf.logging.set_verbosity(tf.logging.INFO) start_time = time.time() # Create the output path. If this fails, it really ought to fail now. :) if not os.path.isdir(FLAGS.output_base_path): os.makedirs(FLAGS.output_base_path) # Create and run model with tf.Graph().as_default(): log("creating the model...") model = SwivelModel(FLAGS) # Create a session for running Ops on the Graph. gpu_opts = {} if FLAGS.per_process_gpu_memory_fraction > 0: gpu_opts["per_process_gpu_memory_fraction"] = \ FLAGS.per_process_gpu_memory_fraction else: gpu_opts["allow_growth"] = True gpu_options = tf.GPUOptions(**gpu_opts) sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) if FLAGS.logs: model.initialize_summary(sess) # Run the Op to initialize the variables. log("initializing the variables...") sess.run(tf.global_variables_initializer()) # Start feeding input log("starting the input threads...") coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess=sess, coord=coord) # Calculate how many steps each thread should run n_total_steps = int(FLAGS.num_epochs * model.n_rows * model.n_cols) / ( FLAGS.submatrix_rows * FLAGS.submatrix_cols) n_steps_per_thread = n_total_steps / ( FLAGS.num_concurrent_steps * model.devices_number) n_submatrices_to_train = model.n_submatrices * FLAGS.num_epochs t0 = [time.time()] n_steps_between_status_updates = 100 n_steps_between_summary_updates = 10000 status_i = [0, 0] status_lock = threading.Lock() msg = ("%%%dd/%%d submatrices trained (%%.1f%%%%), " "%%5.1f submatrices/sec | loss %%f") % \ len(str(n_submatrices_to_train)) def TrainingFn(): for _ in range(int(n_steps_per_thread)): _, global_step, loss = sess.run(( model.train_op, model.global_step, model.loss)) show_status = False update_summary = False with status_lock: new_i = global_step // n_steps_between_status_updates if new_i > status_i[0]: status_i[0] = new_i show_status = True new_i = global_step // n_steps_between_summary_updates if new_i > status_i[1]: status_i[1] = new_i update_summary = True if show_status: elapsed = float(time.time() - t0[0]) log(msg, global_step, n_submatrices_to_train, 100.0 * global_step / n_submatrices_to_train, n_steps_between_status_updates / elapsed, loss) t0[0] = time.time() if update_summary and FLAGS.logs: model.write_summary(sess) # Start training threads train_threads = [] for _ in range(FLAGS.num_concurrent_steps): t = threading.Thread(target=TrainingFn) train_threads.append(t) t.start() # Wait for threads to finish. for t in train_threads: t.join() coord.request_stop() coord.join(threads) # Write out vectors write_embeddings_to_disk(FLAGS, model, sess) # Shutdown sess.close() log("Elapsed: %s", time.time() - start_time) if __name__ == "__main__": tf.app.run()