# Copyright 2016 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.
"""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.
Swivel can be run "stand-alone" or "distributed". The latter involves running
at least one parameter server process, along with one or more worker processes.
"""
from __future__ import division
from __future__ import print_function
import glob
import itertools
import os
import random
import numpy as np
import scipy.stats
import tensorflow as tf
flags = tf.app.flags
flags.DEFINE_string(
'input_base_path', '/tmp/swivel_data',
'Directory containing input shards, vocabularies, and marginals.')
flags.DEFINE_string(
'output_base_path', '/tmp/swivel_data',
'Path where to write the trained embeddings.')
flags.DEFINE_string('eval_base_path', '', 'Path to evaluation data')
# Control for training.
flags.DEFINE_float('num_epochs', 40, 'Number epochs to train')
flags.DEFINE_string('hparams', '', 'Model hyper-parameters')
# Model hyper-parameters. (Move these to tf.HParams once that gets integrated
# into TF from tf.contrib.)
flags.DEFINE_integer(
'dim', 300, 'Embedding dimensionality')
flags.DEFINE_string(
'optimizer', 'rmsprop', 'SGD optimizer; either "adagrad" or "rmsprop"')
flags.DEFINE_float(
'learning_rate', 0.1, 'Optimizer learning rate')
flags.DEFINE_float(
'momentum', 0.1, 'Optimizer momentum; used with RMSProp')
flags.DEFINE_float(
'confidence_base', 0.0, 'Base for count weighting')
flags.DEFINE_float(
'confidence_scale', 1.0, 'Scale for count weighting')
flags.DEFINE_float(
'confidence_exponent', 0.5, 'Exponent for count weighting')
flags.DEFINE_integer(
'submatrix_rows', 4096, 'Number of rows in each submatrix')
flags.DEFINE_integer(
'submatrix_cols', 4096, 'Number of cols in each submatrix')
# For distributed training.
flags.DEFINE_string(
'ps_hosts', '',
'Comma-separated list of parameter server host:port; if empty, run local')
flags.DEFINE_string(
'worker_hosts', '', 'Comma-separated list of worker host:port')
flags.DEFINE_string(
'job_name', '', 'The job this process will run, either "ps" or "worker"')
flags.DEFINE_integer(
'task_index', 0, 'The task index for this process')
flags.DEFINE_integer(
'gpu_device', 0, 'The GPU device to use.')
FLAGS = flags.FLAGS
class Model(object):
"""A Swivel model."""
def __init__(self, input_base_path, hparams):
"""Creates a new Swivel model."""
# Read vocab
self.row_ix_to_word, self.row_word_to_ix = self._read_vocab(
os.path.join(input_base_path, 'row_vocab.txt'))
self.col_ix_to_word, self.col_word_to_ix = self._read_vocab(
os.path.join(input_base_path, 'col_vocab.txt'))
# Read marginals.
row_sums = self._read_marginals_file(
os.path.join(input_base_path, 'row_sums.txt'))
col_sums = self._read_marginals_file(
os.path.join(input_base_path, 'col_sums.txt'))
# Construct input tensors.
count_matrix_files = glob.glob(
os.path.join(input_base_path, 'shard-*.pb'))
global_rows, global_cols, counts = self._count_matrix_input(
count_matrix_files, hparams.submatrix_rows, hparams.submatrix_cols)
# Create embedding variables.
sigma = 1.0 / np.sqrt(hparams.dim)
self.row_embedding = tf.get_variable(
'row_embedding',
shape=[len(row_sums), hparams.dim],
initializer=tf.random_normal_initializer(0, sigma),
dtype=tf.float32)
self.col_embedding = tf.get_variable(
'col_embedding',
shape=[len(col_sums), hparams.dim],
initializer=tf.random_normal_initializer(0, sigma),
dtype=tf.float32)
matrix_log_sum = np.log(np.sum(row_sums) + 1)
row_bias = tf.constant(
[np.log(x + 1) for x in row_sums], dtype=tf.float32)
col_bias = tf.constant(
[np.log(x + 1) for x in col_sums], dtype=tf.float32)
# Fetch embeddings.
selected_rows = tf.nn.embedding_lookup(self.row_embedding, global_rows)
selected_cols = tf.nn.embedding_lookup(self.col_embedding, global_cols)
selected_row_bias = tf.gather(row_bias, global_rows)
selected_col_bias = tf.gather(col_bias, global_cols)
predictions = tf.matmul(selected_rows, selected_cols, transpose_b=True)
# These binary masks separate zero from non-zero values.
count_is_nonzero = tf.to_float(tf.cast(counts, tf.bool))
count_is_zero = 1 - count_is_nonzero
objectives = count_is_nonzero * tf.log(counts + 1e-30)
objectives -= tf.reshape(selected_row_bias, [-1, 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 = (hparams.confidence_base +
hparams.confidence_scale * tf.pow(
counts, hparams.confidence_exponent))
loss_multiplier = 1 / np.sqrt(
hparams.submatrix_rows * hparams.submatrix_cols)
l2_loss = loss_multiplier * tf.reduce_sum(
0.5 * l2_confidence * tf.square(err))
sigmoid_loss = loss_multiplier * tf.reduce_sum(
tf.nn.softplus(err) * count_is_zero)
self.loss_op = l2_loss + sigmoid_loss
if hparams.optimizer == 'adagrad':
opt = tf.train.AdagradOptimizer(hparams.learning_rate)
elif hparams.optimizer == 'rmsprop':
opt = tf.train.RMSPropOptimizer(hparams.learning_rate, hparams.momentum)
else:
raise ValueError('unknown optimizer "%s"' % hparams.optimizer)
self.global_step = tf.get_variable(
'global_step', initializer=0, trainable=False)
self.train_op = opt.minimize(self.loss_op, global_step=self.global_step)
# One epoch trains each submatrix once.
self.steps_per_epoch = (
(len(row_sums) / hparams.submatrix_rows) *
(len(col_sums) / hparams.submatrix_cols))
def _read_vocab(self, filename):
"""Reads the vocabulary file."""
with open(filename) as lines:
ix_to_word = [line.strip() for line in lines]
word_to_ix = {word: ix for ix, word in enumerate(ix_to_word)}
return ix_to_word, word_to_ix
def _read_marginals_file(self, filename):
"""Reads text file with one number per line to an array."""
with open(filename) as lines:
return [float(line.strip()) for line in lines]
def _count_matrix_input(self, filenames, submatrix_rows, submatrix_cols):
"""Creates ops that read submatrix shards from disk."""
random.shuffle(filenames)
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)
return global_row, global_col, count
def wordsim_eval_op(self, filename):
"""Returns an op that runs an eval on a word similarity dataset.
The eval dataset is assumed to be tab-separated, one scored word pair per
line. The resulting value is Spearman's rho of the human judgements with
the cosine similarity of the word embeddings.
Args:
filename: the filename containing the word similarity data.
Returns:
An operator that will compute Spearman's rho of the current row
embeddings.
"""
with open(filename, 'r') as fh:
tuples = (line.strip().split('\t') for line in fh.read().splitlines())
word1s, word2s, sims = zip(*tuples)
actuals = map(float, sims)
v1s_t = tf.nn.embedding_lookup(
self.row_embedding,
[self.row_word_to_ix.get(w, 0) for w in word1s])
v2s_t = tf.nn.embedding_lookup(
self.row_embedding,
[self.row_word_to_ix.get(w, 0) for w in word2s])
# Compute the predicted word similarity as the cosine similarity between the
# embedding vectors.
preds_t = tf.reduce_sum(
tf.nn.l2_normalize(v1s_t, dim=1) * tf.nn.l2_normalize(v2s_t, dim=1),
axis=1)
def _op(preds):
rho, _ = scipy.stats.spearmanr(preds, actuals)
return rho
return tf.py_func(_op, [preds_t], tf.float64)
def analogy_eval_op(self, filename, max_vocab_size=20000):
"""Returns an op that runs an eval on an analogy dataset.
The eval dataset is assumed to be tab-separated, with four tokens per
line. The first three tokens are query terms, the last is the expected
answer. For each line (e.g., "man king woman queen"), the vectors
corresponding to the query terms are added ("king - man + woman") to produce
a query vector. If the expected answer's vector is the nearest neighbor to
the query vector (not counting any of the query vectors themselves), then
the line is scored as correct. The reported accuracy is the number of
correct rows divided by the total number of rows. Missing terms are
replaced with an arbitrary vector and will almost certainly result in
incorrect answers.
Note that the results are approximate: for efficiency's sake, only the first
`max_vocab_size` terms are included in the nearest neighbor search.
Args:
filename: the filename containing the analogy data.
max_vocab_size: the maximum number of tokens to include in the nearest
neighbor search. By default, 20000.
Returns:
The accuracy on the analogy task.
"""
analogy_ixs = []
with open(filename, 'r') as lines:
for line in lines:
parts = line.strip().split('\t')
if len(parts) == 4:
analogy_ixs.append([self.row_word_to_ix.get(w, 0) for w in parts])
# man:king :: woman:queen => king - man + woman == queen
ix1s, ix2s, ix3s, _ = zip(*analogy_ixs)
v1s_t, v2s_t, v3s_t = (
tf.nn.l2_normalize(
tf.nn.embedding_lookup(self.row_embedding, ixs),
dim=1)
for ixs in (ix1s, ix2s, ix3s))
preds_t = v2s_t - v1s_t + v3s_t
# Compute the nearest neighbors as the cosine similarity. We only consider
# up to max_vocab_size to avoid a matmul that swamps the machine.
sims_t = tf.matmul(
preds_t,
tf.nn.l2_normalize(self.row_embedding[:max_vocab_size], dim=1),
transpose_b=True)
# Take the four nearest neighbors, since the eval explicitly discards the
# query terms.
_, preds_ixs_t = tf.nn.top_k(sims_t, 4)
def _op(preds_ixs):
correct, total = 0, 0
for pred_ixs, actual_ixs in itertools.izip(preds_ixs, analogy_ixs):
pred_ixs = [ix for ix in pred_ixs if ix not in actual_ixs[:3]]
correct += pred_ixs[0] == actual_ixs[3]
total += 1
return correct / total
return tf.py_func(_op, [preds_ixs_t], tf.float64)
def _write_tensor(self, vocab_path, output_path, session, embedding):
"""Writes tensor to output_path as tsv."""
embeddings = session.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]
print('\t'.join([word.strip()] + [str(x) for x in embedding]),
file=out_f)
def write_embeddings(self, config, session):
"""Writes row and column embeddings disk."""
self._write_tensor(
os.path.join(config.input_base_path, 'row_vocab.txt'),
os.path.join(config.output_base_path, 'row_embedding.tsv'),
session, self.row_embedding)
self._write_tensor(
os.path.join(config.input_base_path, 'col_vocab.txt'),
os.path.join(config.output_base_path, 'col_embedding.tsv'),
session, self.col_embedding)
def main(_):
tf.logging.set_verbosity(tf.logging.INFO)
# If we have ps_hosts, then we'll assume that this is going to be a
# distributed training run. Configure the cluster appropriately. Otherwise,
# we just do everything in-process.
if FLAGS.ps_hosts:
cluster = tf.train.ClusterSpec({
'ps': FLAGS.ps_hosts.split(','),
'worker': FLAGS.worker_hosts.split(','),
})
if FLAGS.job_name == 'ps':
# Ignore the GPU if we're the parameter server. This let's the PS run on
# the same machine as a worker.
config = tf.ConfigProto(device_count={'GPU': 0})
elif FLAGS.job_name == 'worker':
config = tf.ConfigProto(gpu_options=tf.GPUOptions(
visible_device_list='%d' % FLAGS.gpu_device,
allow_growth=True))
else:
raise ValueError('unknown job name "%s"' % FLAGS.job_name)
server = tf.train.Server(
cluster,
job_name=FLAGS.job_name,
task_index=FLAGS.task_index,
config=config)
if FLAGS.job_name == 'ps':
return server.join()
device_setter = tf.train.replica_device_setter(
worker_device='/job:worker/task:%d' % FLAGS.task_index,
cluster=cluster)
else:
server = None
device_setter = tf.train.replica_device_setter(0)
# Build the graph.
with tf.Graph().as_default():
with tf.device(device_setter):
model = Model(FLAGS.input_base_path, FLAGS)
# If an eval path is present, then create eval operators and set up scalar
# summaries to report on the results. Run the evals on the CPU since
# the analogy eval requires a fairly enormous tensor to be allocated to
# do the nearest neighbor search.
if FLAGS.eval_base_path:
wordsim_filenames = glob.glob(
os.path.join(FLAGS.eval_base_path, '*.ws.tab'))
for filename in wordsim_filenames:
name = os.path.basename(filename).split('.')[0]
with tf.device(tf.DeviceSpec(device_type='CPU')):
op = model.wordsim_eval_op(filename)
tf.summary.scalar(name, op)
analogy_filenames = glob.glob(
os.path.join(FLAGS.eval_base_path, '*.an.tab'))
for filename in analogy_filenames:
name = os.path.basename(filename).split('.')[0]
with tf.device(tf.DeviceSpec(device_type='CPU')):
op = model.analogy_eval_op(filename)
tf.summary.scalar(name, op)
tf.summary.scalar('loss', model.loss_op)
# Train on, soldier.
supervisor = tf.train.Supervisor(
logdir=FLAGS.output_base_path,
is_chief=(FLAGS.task_index == 0),
save_summaries_secs=60,
recovery_wait_secs=5)
max_step = FLAGS.num_epochs * model.steps_per_epoch
master = server.target if server else ''
with supervisor.managed_session(master) as session:
local_step = 0
global_step = session.run(model.global_step)
while not supervisor.should_stop() and global_step < max_step:
global_step, loss, _ = session.run([
model.global_step, model.loss_op, model.train_op])
if not np.isfinite(loss):
raise ValueError('non-finite cost at step %d' % global_step)
local_step += 1
if local_step % 10 == 0:
tf.logging.info(
'local_step=%d global_step=%d loss=%.1f, %.1f%% complete',
local_step, global_step, loss, 100.0 * global_step / max_step)
if FLAGS.task_index == 0:
supervisor.saver.save(
session, supervisor.save_path, global_step=global_step)
model.write_embeddings(FLAGS, session)
if __name__ == '__main__':
tf.app.run()
