def main(_):
opts = Options(save_path=FLAGS.save_path,
train_biom=FLAGS.train_biom,
test_biom=FLAGS.test_biom,
train_metadata=FLAGS.train_metadata,
test_metadata=FLAGS.test_metadata,
formula=FLAGS.formula,
learning_rate=FLAGS.learning_rate,
clipping_size=FLAGS.clipping_size,
beta_mean=FLAGS.beta_mean,
beta_scale=FLAGS.beta_scale,
gamma_mean=FLAGS.gamma_mean,
gamma_scale=FLAGS.gamma_scale,
epochs_to_train=FLAGS.epochs_to_train,
num_neg_samples=FLAGS.num_neg_samples,
batch_size=FLAGS.batch_size,
min_sample_count=FLAGS.min_sample_count,
min_feature_count=FLAGS.min_feature_count,
statistics_interval=FLAGS.statistics_interval,
summary_interval=FLAGS.summary_interval,
checkpoint_interval=FLAGS.checkpoint_interval)
# preprocessing
train_table, train_metadata = opts.train_table, opts.train_metadata
train_metadata = train_metadata.loc[train_table.ids(axis='sample')]
sample_filter = lambda val, id_, md: (
(id_ in train_metadata.index) and np.sum(val) > opts.min_sample_count)
read_filter = lambda val, id_, md: np.sum(val) > opts.min_feature_count
metadata_filter = lambda val, id_, md: id_ in train_metadata.index
train_table = train_table.filter(metadata_filter, axis='sample')
train_table = train_table.filter(sample_filter, axis='sample')
train_table = train_table.filter(read_filter, axis='observation')
train_metadata = train_metadata.loc[train_table.ids(axis='sample')]
sort_f = lambda xs: [xs[train_metadata.index.get_loc(x)] for x in xs]
train_table = train_table.sort(sort_f=sort_f, axis='sample')
train_metadata = dmatrix(opts.formula,
train_metadata,
return_type='dataframe')
# hold out data preprocessing
test_table, test_metadata = opts.test_table, opts.test_metadata
metadata_filter = lambda val, id_, md: id_ in test_metadata.index
obs_lookup = set(train_table.ids(axis='observation'))
feat_filter = lambda val, id_, md: id_ in obs_lookup
test_table = test_table.filter(metadata_filter, axis='sample')
test_table = test_table.filter(feat_filter, axis='observation')
test_metadata = test_metadata.loc[test_table.ids(axis='sample')]
sort_f = lambda xs: [xs[test_metadata.index.get_loc(x)] for x in xs]
test_table = test_table.sort(sort_f=sort_f, axis='sample')
test_metadata = dmatrix(opts.formula,
test_metadata,
return_type='dataframe')
p = train_metadata.shape[1] # number of covariates
G_data = train_metadata.values
y_data = np.array(train_table.matrix_data.todense()).T
y_test = np.array(test_table.matrix_data.todense()).T
N, D = y_data.shape
save_path = opts.save_path
learning_rate = opts.learning_rate
batch_size = opts.batch_size
gamma_mean, gamma_scale = opts.gamma_mean, opts.gamma_scale
beta_mean, beta_scale = opts.beta_mean, opts.beta_scale
num_iter = (N // batch_size) * opts.epochs_to_train
holdout_size = test_metadata.shape[0]
checkpoint_interval = opts.checkpoint_interval
# Model code
with tf.Graph().as_default(), tf.Session() as session:
with tf.device("/cpu:0"):
# Place holder variables to accept input data
G_ph = tf.placeholder(tf.float32, [batch_size, p], name='G_ph')
Y_ph = tf.placeholder(tf.float32, [batch_size, D], name='Y_ph')
G_holdout = tf.placeholder(tf.float32, [holdout_size, p],
name='G_holdout')
Y_holdout = tf.placeholder(tf.float32, [holdout_size, D],
name='Y_holdout')
total_count = tf.placeholder(tf.float32, [batch_size],
name='total_count')
# Define PointMass Variables first
qgamma = tf.Variable(tf.random_normal([1, D]), name='qgamma')
qbeta = tf.Variable(tf.random_normal([p, D]), name='qB')
# Distributions
# species bias
gamma = Normal(loc=tf.zeros([1, D]) + gamma_mean,
scale=tf.ones([1, D]) * gamma_scale,
name='gamma')
# regression coefficents distribution
beta = Normal(loc=tf.zeros([p, D]) + beta_mean,
scale=tf.ones([p, D]) * beta_scale,
name='B')
Bprime = tf.concat([qgamma, qbeta], axis=0)
# add bias terms for samples
Gprime = tf.concat([tf.ones([batch_size, 1]), G_ph], axis=1)
eta = tf.matmul(Gprime, Bprime)
phi = tf.nn.log_softmax(eta)
Y = Multinomial(total_count=total_count, logits=phi, name='Y')
loss = -(tf.reduce_mean(gamma.log_prob(qgamma)) + \
tf.reduce_mean(beta.log_prob(qbeta)) + \
tf.reduce_mean(Y.log_prob(Y_ph)) * (N / batch_size))
loss = tf.Print(loss, [loss])
optimizer = tf.train.AdamOptimizer(learning_rate)
gradients, variables = zip(*optimizer.compute_gradients(loss))
gradients, _ = tf.clip_by_global_norm(gradients,
opts.clipping_size)
train = optimizer.apply_gradients(zip(gradients, variables))
with tf.name_scope('accuracy'):
holdout_count = tf.reduce_sum(Y_holdout, axis=1)
pred = tf.reshape(holdout_count, [-1, 1]) * tf.nn.softmax(
tf.matmul(G_holdout, qbeta) + qgamma)
mse = tf.reduce_mean(tf.squeeze(tf.abs(pred - Y_holdout)))
tf.summary.scalar('mean_absolute_error', mse)
tf.summary.scalar('loss', loss)
tf.summary.histogram('qbeta', qbeta)
tf.summary.histogram('qgamma', qgamma)
merged = tf.summary.merge_all()
tf.global_variables_initializer().run()
writer = tf.summary.FileWriter(save_path, session.graph)
losses = np.array([0.] * num_iter)
idx = np.arange(train_metadata.shape[0])
log_handle = open(os.path.join(save_path, 'run.log'), 'w')
last_checkpoint_time = 0
start_time = time.time()
saver = tf.train.Saver()
for i in range(num_iter):
batch_idx = np.random.choice(idx, size=batch_size)
feed_dict = {
Y_ph: y_data[batch_idx].astype(np.float32),
G_ph: train_metadata.values[batch_idx].astype(np.float32),
Y_holdout: y_test.astype(np.float32),
G_holdout: test_metadata.values.astype(np.float32),
total_count:
y_data[batch_idx].sum(axis=1).astype(np.float32)
}
if i % 1000 == 0:
run_options = tf.RunOptions(
trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
_, summary, train_loss, grads = session.run(
[train, merged, loss, gradients],
feed_dict=feed_dict,
options=run_options,
run_metadata=run_metadata)
writer.add_run_metadata(run_metadata, 'step%d' % i)
writer.add_summary(summary, i)
elif i % 5000 == 0:
_, summary, err, train_loss, grads = session.run(
[train, mse, merged, loss, gradients],
feed_dict=feed_dict)
writer.add_summary(summary, i)
else:
_, summary, train_loss, grads = session.run(
[train, merged, loss, gradients], feed_dict=feed_dict)
writer.add_summary(summary, i)
now = time.time()
if now - last_checkpoint_time > checkpoint_interval:
saver.save(session,
os.path.join(opts.save_path, "model.ckpt"),
global_step=i)
last_checkpoint_time = now
losses[i] = train_loss
elapsed_time = time.time() - start_time
print('Elapsed Time: %f seconds' % elapsed_time)
# Cross validation
pred_beta = qbeta.eval()
pred_gamma = qgamma.eval()
mse, mrc = cross_validation(test_metadata.values, pred_beta,
pred_gamma, y_test)
print("MSE: %f, MRC: %f" % (mse, mrc))
def __call__(self, session, trainX, trainY, testX, testY):
""" Initialize the actual graph
Parameters
----------
session : tf.Session
Tensorflow session
trainX : sparse array in coo format
Test input OTU table, where rows are samples and columns are
observations
trainY : np.array
Test output metabolite table
testX : sparse array in coo format
Test input OTU table, where rows are samples and columns are
observations. This is mainly for cross validation.
testY : np.array
Test output metabolite table. This is mainly for cross validation.
"""
self.session = session
self.nnz = len(trainX.data)
self.d1 = trainX.shape[1]
self.d2 = trainY.shape[1]
self.cv_size = len(testX.data)
# keep the multinomial sampling on the cpu
# https://github.com/tensorflow/tensorflow/issues/18058
with tf.device('/cpu:0'):
X_ph = tf.SparseTensor(indices=np.array([trainX.row,
trainX.col]).T,
values=trainX.data,
dense_shape=trainX.shape)
Y_ph = tf.constant(trainY, dtype=tf.float32)
X_holdout = tf.SparseTensor(indices=np.array(
[testX.row, testX.col]).T,
values=testX.data,
dense_shape=testX.shape)
Y_holdout = tf.constant(testY, dtype=tf.float32)
total_count = tf.reduce_sum(Y_ph, axis=1)
batch_ids = tf.multinomial(
tf.log(tf.reshape(X_ph.values, [1, -1])), self.batch_size)
batch_ids = tf.squeeze(batch_ids)
X_samples = tf.gather(X_ph.indices, 0, axis=1)
X_obs = tf.gather(X_ph.indices, 1, axis=1)
sample_ids = tf.gather(X_samples, batch_ids)
Y_batch = tf.gather(Y_ph, sample_ids)
X_batch = tf.gather(X_obs, batch_ids)
with tf.device(self.device_name):
self.qUmain = tf.Variable(tf.random_normal([self.d1, self.p]),
name='qU')
self.qUbias = tf.Variable(tf.random_normal([self.d1, 1]),
name='qUbias')
self.qVmain = tf.Variable(tf.random_normal([self.p, self.d2 - 1]),
name='qV')
self.qVbias = tf.Variable(tf.random_normal([1, self.d2 - 1]),
name='qVbias')
qU = tf.concat([tf.ones([self.d1, 1]), self.qUbias, self.qUmain],
axis=1)
qV = tf.concat(
[self.qVbias,
tf.ones([1, self.d2 - 1]), self.qVmain], axis=0)
# regression coefficents distribution
Umain = Normal(loc=tf.zeros([self.d1, self.p]) + self.u_mean,
scale=tf.ones([self.d1, self.p]) * self.u_scale,
name='U')
Ubias = Normal(loc=tf.zeros([self.d1, 1]) + self.u_mean,
scale=tf.ones([self.d1, 1]) * self.u_scale,
name='biasU')
Vmain = Normal(loc=tf.zeros([self.p, self.d2 - 1]) + self.v_mean,
scale=tf.ones([self.p, self.d2 - 1]) * self.v_scale,
name='V')
Vbias = Normal(loc=tf.zeros([1, self.d2 - 1]) + self.v_mean,
scale=tf.ones([1, self.d2 - 1]) * self.v_scale,
name='biasV')
du = tf.gather(qU, X_batch, axis=0, name='du')
dv = tf.concat([tf.zeros([self.batch_size, 1]), du @ qV],
axis=1,
name='dv')
tc = tf.gather(total_count, sample_ids)
Y = Multinomial(total_count=tc, logits=dv, name='Y')
num_samples = trainX.shape[0]
norm = num_samples / self.batch_size
logprob_vmain = tf.reduce_sum(Vmain.log_prob(self.qVmain),
name='logprob_vmain')
logprob_vbias = tf.reduce_sum(Vbias.log_prob(self.qVbias),
name='logprob_vbias')
logprob_umain = tf.reduce_sum(Umain.log_prob(self.qUmain),
name='logprob_umain')
logprob_ubias = tf.reduce_sum(Ubias.log_prob(self.qUbias),
name='logprob_ubias')
logprob_y = tf.reduce_sum(Y.log_prob(Y_batch), name='logprob_y')
self.log_loss = -(logprob_y * norm + logprob_umain +
logprob_ubias + logprob_vmain + logprob_vbias)
# keep the multinomial sampling on the cpu
# https://github.com/tensorflow/tensorflow/issues/18058
with tf.device('/cpu:0'):
# cross validation
with tf.name_scope('accuracy'):
cv_batch_ids = tf.multinomial(
tf.log(tf.reshape(X_holdout.values, [1, -1])),
self.cv_size)
cv_batch_ids = tf.squeeze(cv_batch_ids)
X_cv_samples = tf.gather(X_holdout.indices, 0, axis=1)
X_cv = tf.gather(X_holdout.indices, 1, axis=1)
cv_sample_ids = tf.gather(X_cv_samples, cv_batch_ids)
Y_cvbatch = tf.gather(Y_holdout, cv_sample_ids)
X_cvbatch = tf.gather(X_cv, cv_batch_ids)
holdout_count = tf.reduce_sum(Y_cvbatch, axis=1)
cv_du = tf.gather(qU, X_cvbatch, axis=0, name='cv_du')
pred = tf.reshape(holdout_count, [-1, 1]) * tf.nn.softmax(
tf.concat([tf.zeros([self.cv_size, 1]), cv_du @ qV],
axis=1,
name='pred'))
self.cv = tf.reduce_mean(tf.squeeze(tf.abs(pred - Y_cvbatch)))
# keep all summaries on the cpu
with tf.device('/cpu:0'):
tf.summary.scalar('logloss', self.log_loss)
tf.summary.scalar('cv_rmse', self.cv)
tf.summary.histogram('qUmain', self.qUmain)
tf.summary.histogram('qVmain', self.qVmain)
tf.summary.histogram('qUbias', self.qUbias)
tf.summary.histogram('qVbias', self.qVbias)
self.merged = tf.summary.merge_all()
self.writer = tf.summary.FileWriter(self.save_path,
self.session.graph)
with tf.device(self.device_name):
with tf.name_scope('optimize'):
optimizer = tf.train.AdamOptimizer(self.learning_rate,
beta1=self.beta_1,
beta2=self.beta_2)
gradients, self.variables = zip(
*optimizer.compute_gradients(self.log_loss))
self.gradients, _ = tf.clip_by_global_norm(
gradients, self.clipnorm)
self.train = optimizer.apply_gradients(
zip(self.gradients, self.variables))
tf.global_variables_initializer().run()
def __call__(self, session, trainX, trainY, testX, testY):
""" Initialize the actual graph
Parameters
----------
session : tf.Session
Tensorflow session
trainX : np.array
Input training design matrix.
trainY : np.array
Output training OTU table, where rows are samples and columns are
observations.
testX : np.array
Input testing design matrix.
testY : np.array
Output testing OTU table, where rows are samples and columns are
observations.
"""
self.session = session
self.N, self.p = trainX.shape
self.D = trainY.shape[1]
holdout_size = testX.shape[0]
# Place holder variables to accept input data
self.X_ph = tf.constant(trainX, dtype=tf.float32, name='G_ph')
self.Y_ph = tf.constant(trainY, dtype=tf.float32, name='Y_ph')
self.X_holdout = tf.constant(testX, dtype=tf.float32, name='G_holdout')
self.Y_holdout = tf.constant(testY, dtype=tf.float32, name='Y_holdout')
batch_ids = tf.multinomial(tf.ones([1, self.N]), self.batch_size)
sample_ids = tf.squeeze(batch_ids)
Y_batch = tf.gather(self.Y_ph, sample_ids, axis=0)
X_batch = tf.gather(self.X_ph, sample_ids, axis=0)
total_count = tf.reduce_sum(Y_batch, axis=1)
holdout_count = tf.reduce_sum(self.Y_holdout, axis=1)
# Define PointMass Variables first
self.qbeta = tf.Variable(tf.random_normal([self.p, self.D - 1]),
name='qB')
# regression coefficents distribution
beta = Normal(loc=tf.zeros([self.p, self.D - 1]) + self.beta_mean,
scale=tf.ones([self.p, self.D - 1]) * self.beta_scale,
name='B')
eta = tf.matmul(X_batch, self.qbeta, name='eta')
phi = tf.nn.log_softmax(tf.concat(
[tf.zeros([self.batch_size, 1]), eta], axis=1),
name='phi')
Y = Multinomial(total_count=total_count, logits=phi, name='Y')
# cross validation
with tf.name_scope('accuracy'):
pred = tf.reshape(holdout_count, [-1, 1]) * tf.nn.softmax(
tf.concat([
tf.zeros([holdout_size, 1]),
tf.matmul(self.X_holdout, self.qbeta)
],
axis=1),
name='phi')
self.cv = tf.reduce_mean(tf.squeeze(tf.abs(pred - self.Y_holdout)))
tf.summary.scalar('mean_absolute_error', self.cv)
self.loss = -(tf.reduce_sum(beta.log_prob(self.qbeta)) +
tf.reduce_sum(Y.log_prob(Y_batch)) *
(self.N / self.batch_size))
optimizer = tf.train.AdamOptimizer(self.learning_rate,
beta1=self.beta_1,
beta2=self.beta_2)
gradients, variables = zip(*optimizer.compute_gradients(self.loss))
self.gradients, _ = tf.clip_by_global_norm(gradients, self.clipnorm)
self.train = optimizer.apply_gradients(zip(gradients, variables))
tf.summary.scalar('loss', self.loss)
tf.summary.histogram('qbeta', self.qbeta)
self.merged = tf.summary.merge_all()
if self.save_path is not None:
self.writer = tf.summary.FileWriter(self.save_path,
self.session.graph)
else:
self.writer = None
tf.global_variables_initializer().run()