""" Code for the Main function of MetaPred. """
import os, csv
import numpy as np
import random
import pickle as pkl
import tensorflow as tf
import copy
from tensorflow.python.platform import flags
from data_loader import DataLoader
import model, finetune
FLAGS = flags.FLAGS
flags.DEFINE_string('source', 'AD', 'source task')
flags.DEFINE_string('target', 'MCI', 'simulated task')
flags.DEFINE_string('true_target', 'PD', 'true task')
## Dataset/method options
flags.DEFINE_integer('n_classes', 2, 'number of classes used in classification (e.g. binary classification)')
## Training options
flags.DEFINE_string('method', 'rnn', 'deep learning methods for modeling')
flags.DEFINE_integer('pretrain_iterations', 20000, 'number of pre-training iterations')
flags.DEFINE_integer('metatrain_iterations', 10000, 'number of metatraining iterations') # 15k for omniglot, 50k for sinusoid
flags.DEFINE_integer('meta_batch_size', 8, 'number of tasks sampled per meta-update')
flags.DEFINE_integer('update_batch_size', 16, 'number of samples used for inner gradient update (K for K-shot learning)')
flags.DEFINE_float('meta_lr', 0.0001, 'the base learning rate of the generator')
flags.DEFINE_float('update_lr', 1e-3, 'step size alpha for inner gradient update')
flags.DEFINE_integer('num_updates', 4, 'number of inner gradient updates during training')
flags.DEFINE_integer('n_total_batches', 100000, 'total batches generated by random sampling')
## Model options
flags.DEFINE_string('norm', 'None', 'batch_norm, layer_norm, or None')
flags.DEFINE_bool('stop_grad', False, 'if True, do not use second derivatives in meta-optimization (for speed)')
flags.DEFINE_bool('isReg', True, 'if True, compute regularization of weights and bias')
flags.DEFINE_float('dropout', 0.5, 'drop out when modeling, with probability keep_prob')
## Logging, saving, and testing options
flags.DEFINE_integer('run_time', 1, 're-run for stable analysis')
flags.DEFINE_bool('train', True, 'True to train, False to test directly')
flags.DEFINE_bool('test', True, 'True to test, no matter the model is trained')
flags.DEFINE_bool('finetune', False, 'True to finetunning furthermore, after meta-learning')
flags.DEFINE_bool('log', True, 'if false, do not log summaries, for debugging code')
flags.DEFINE_string('logdir', 'model/', 'directory for summaries and checkpoints')
flags.DEFINE_bool('resume', False, 'resume training if there is a model available')
flags.DEFINE_integer('test_iter', -1, 'iteration to load model (-1 for latest model)')
flags.DEFINE_integer('train_update_batch_size', -1, 'number of examples used for gradient update during training (use if you want to test with a different number)')
flags.DEFINE_float('train_update_lr', -1, 'value of inner gradient step step during training. (use if you want to test with a different value)') # 0.1 for omniglot
def train(data_loader, ifold, exp_string):
# construct MetaPred model
print ("constructing MetaPred model ...")
m1 = model.MetaPred(data_loader, FLAGS.meta_lr, FLAGS.update_lr)
# fitting the meta-learning model
print ("model training...")
sess = m1.fit(data_loader.episode, data_loader.episode_val[ifold], ifold, exp_string)
return m1, sess
def test(data_loader, ifold, m, sess, exp_string):
# meta-testing the model
print ("model test...")
data_tuple_val = (data_loader.data_s, data_loader.data_tt_val[ifold], data_loader.label_s, data_loader.label_tt_val[ifold])
test_accs, test_aucs, test_ap, test_f1s = m.evaluate(data_loader.episode_val[ifold], data_tuple_val, sess=sess, prefix="metatest_")
print('Test results: ' + "ifold: " + str(ifold) + ": tAcc: " + str(test_accs) + \
", tAuc: " + str(test_aucs) + ", tAP: " + str(test_ap) + ", tF1: " + str(test_f1s))
return test_accs, test_aucs, test_ap, test_f1s
def fine_tune(data_loader, ifold, meta_m, weights_for_finetune, exp_string):
# construct MetaPred model
is_finetune = True
print ("finetunning MetaPred model ...")
if FLAGS.method == "cnn":
m2 = finetune.CNN(data_loader, weights_for_finetune, freeze_opt=freeze_opt, is_finetune=is_finetune)
if FLAGS.method == "rnn":
m2 = finetune.RNN(data_loader, weights_for_finetune, freeze_opt=freeze_opt, is_finetune=is_finetune)
print ("model finetunning...")
# model finetunning
sess, _, _ = m2.fit(data_loader.tt_sample[ifold], data_loader.tt_label[ifold],
data_loader.tt_sample_val[ifold], data_loader.tt_label_val[ifold])
return m2, sess
def save_results(metatest, exp_string):
out_filename = "results/res_" + exp_string
with open(out_filename, 'w') as f:
writer = csv.writer(f, delimiter=',')
for key in metatest:
writer.writerow([np.mean(np.array(metatest[key]))])
writer.writerow([np.std(np.array(metatest[key]))])
print ("results saved")
def save_weights(meta_m, source, target, true_target, data_loader, ifold):
with open("weights/meta-" + FLAGS.method + ".weights" + ".source_" + "-".join(source) + ".starget_" + "".join(target) + ".ttarget_" + "".join(true_target) + ".pkl", 'wb') as f:
pkl.dump((meta_m.weights_for_finetune), f, protocol=2)
f.close()
with open("weights/meta-" + FLAGS.method + ".tt_train" + ".source_" + "-".join(source) + ".starget_" + "".join(target) + ".ttarget_" + "".join(true_target) + ".pkl", 'wb') as f:
pkl.dump((data_loader.tt_sample[ifold], data_loader.tt_label[ifold]), f, protocol=2)
f.close()
with open("weights/meta-" + FLAGS.method + ".tt_val" + ".source_" + "-".join(source) + ".starget_" + "".join(target) + ".ttarget_" + "".join(true_target) + ".pkl", 'wb') as f:
pkl.dump((data_loader.tt_sample_val[ifold], data_loader.tt_label_val[ifold]), f, protocol=2)
f.close()
print("model weights saved")
def main():
print (FLAGS.method)
# set source and simulated target for training
print ('task setting: ')
source = [FLAGS.source]
target = [FLAGS.target]
true_target = [FLAGS.true_target]
print ("The applied source tasks are: ", " ".join(source))
print ("The simulated target task is: ", " ".join(target))
print ("The true target task is: ", " ".join(true_target))
n_tasks = len(source) + len(target)
# load ehrs data
data_loader = DataLoader(source, target, true_target, n_tasks,
FLAGS.update_batch_size, FLAGS.meta_batch_size)
exp_string = 'stsk_'+str('&'.join(source))+'ttsk_'+str('&'.join(target))+'.mbs_'+str(FLAGS.meta_batch_size) + \
'.ubs_' + str(FLAGS.update_batch_size) + '.numstep' + str(FLAGS.num_updates) + '.updatelr' + str(FLAGS.update_lr)
metatest = {'aucroc': [], 'avepre': [], 'f1score': []} # n_fold result
n_fold = data_loader.n_fold
for ifold in range(n_fold):
print ("----------The %d-th fold-----------" %(ifold+1))
meta_model = None
if FLAGS.train:
meta_model, sess = train(data_loader, ifold, exp_string)
save_weights(meta_model, source, target, true_target, data_loader, ifold)
if FLAGS.finetune:
with open("weights/meta-" + FLAGS.method + ".weights" + ".source_" + "-".join(source) + ".starget_" + "".join(target) + ".ttarget_" + "".join(true_target) + ".pkl", 'rb') as f:
weights_for_finetune = pkl.load(f)
f.close()
model, sess = fine_tune(data_loader, ifold, meta_model, weights_for_finetune, exp_string)
if FLAGS.test:
_, test_aucs, test_ap, test_f1s = test(data_loader, ifold, meta_model, sess, exp_string)
metatest['aucroc'].append(test_aucs)
metatest['avepre'].append(test_ap)
metatest['f1score'].append(test_f1s)
# show results
print ('--------------- model setting ---------------')
print('source: ', " ".join(source), 'simulated target: ', " ".join(target), 'true target: ', " ".join(true_target))
print('method:', 'meta-' + FLAGS.method, 'meta-bz:', FLAGS.meta_batch_size, 'update-bz:', FLAGS.update_batch_size, \
'num update:', FLAGS.num_updates, 'meta-lr:', FLAGS.meta_lr, 'update-lr:', FLAGS.update_lr)
print ('--------------- 5fold results ---------------')
print ('aucroc mean: ', np.mean(np.array(metatest['aucroc'])))
print ('aucroc std: ', np.std(np.array(metatest['aucroc'])))
print ('f1score mean: ', np.mean(np.array(metatest['f1score'])))
print ('f1score std: ', np.std(np.array(metatest['f1score'])))
save_results(metatest, exp_string)
if __name__ == "__main__":
main()
