main_paper.py

"""
Usage Instructions:

# R2D2 exact paper
### miniImagenet
#### 5 way, 1 shot, backprop everywhere, correct dropout RESULT: 51.7 (20k)
python main_paper.py --datasource=miniimagenet --metatrain_iterations=20000 --meta_batch_size=4 --update_batch_size=1 --update_lr=0.01 --meta_lr=0.005 --num_updates=1 --num_classes=5 --logdir=logs/paperFullBPminiimagenet5way1shot/ --num_filters=32 --max_pool=True  

#### 5 way, 5 shot, backprop everywhere, correct dropout RESULT: 64.1 (16k)
python main_paper.py --datasource=miniimagenet --metatrain_iterations=20000 --meta_batch_size=4 --update_batch_size=5 --update_lr=0.01 --meta_lr=0.005 --num_updates=1 --num_classes=5 --logdir=logs/paperFullBPminiimagenet5way5shot/ --num_filters=32 --max_pool=True  

#### 2 way, 1 shot, backprop everywhere, correct dropout RESULT: 74.6 ± 2.9%
python main_paper.py --datasource=miniimagenet --metatrain_iterations=20000 --meta_batch_size=4 --update_batch_size=1 --update_lr=0.01 --meta_lr=0.005 --num_updates=1 --num_classes=2 --logdir=logs/paperFullBPminiimagenet2way1shot/ --num_filters=32 --max_pool=True  

#### 2 way, 5 shot, backprop everywhere, correct dropout RESULT: 
python main_paper.py --datasource=miniimagenet --metatrain_iterations=20000 --meta_batch_size=4 --update_batch_size=5 --update_lr=0.01 --meta_lr=0.005 --num_updates=1 --num_classes=2 --logdir=logs/paperFullBPminiimagenet2way5shot/ --num_filters=32 --max_pool=True  

### CIFAR FS
#### 5 way, 1 shot, RESULT: 
python main_paper.py --datasource=cifarfs --metatrain_iterations=60000 --meta_batch_size=4 --update_batch_size=1 --update_lr=0.01 --meta_lr=0.005 --num_updates=1 --num_classes=5 --logdir=logs/paperFullBPcifarfs5way1shot/ --num_filters=32 --max_pool=True  

#### 2 way, 1 shot, RESULT: 83.6 (20k)
python main_paper.py --datasource=cifarfs --metatrain_iterations=60000 --meta_batch_size=4 --update_batch_size=1 --update_lr=0.01 --meta_lr=0.005 --num_updates=1 --num_classes=2 --logdir=logs/paperFullBPcifarfs2way1shot/ --num_filters=32 --max_pool=True  

#### 2 way, 5 shot, RESULT: 89.0 (20k)
python main_paper.py --datasource=cifarfs --metatrain_iterations=30000 --meta_batch_size=4 --update_batch_size=5 --update_lr=0.01 --meta_lr=0.005 --num_updates=1 --num_classes=2 --logdir=logs/paperFullBPcifarfs2way5shot/ --num_filters=32 --max_pool=True  

#### 5 way, 5 shot, RESULT:
python main_paper.py --datasource=cifarfs --metatrain_iterations=30000 --meta_batch_size=4 --update_batch_size=5 --update_lr=0.01 --meta_lr=0.005 --num_updates=1 --num_classes=5 --logdir=logs/paperFullBPcifarfs5way5shot/ --num_filters=32 --max_pool=True 
        
    To run evaluation, use the '--train=False' flag and the '--test_set=True' flag to use the test set.

    For miniimagenet training, acquire the dataset online, put it in the correspoding data directory, and see the python script instructions in that directory to preprocess the data. For CIFAR fs training, the dataset is automatically downloaded, and the splits are present in the code in the data directory.
"""
import csv
import numpy as np
import pickle
import random
import tensorflow as tf

from data_generator import DataGenerator
from r2d2_paper import R2D2_paper
from tensorflow.python.platform import flags

FLAGS = flags.FLAGS

## Dataset/method options
flags.DEFINE_string('datasource', 'sinusoid', 'sinusoid or omniglot or miniimagenet')
flags.DEFINE_integer('num_classes', 5, 'number of classes used in classification (e.g. 5-way classification).')
# oracle means task id is input (only suitable for sinusoid)
flags.DEFINE_string('baseline', None, 'oracle, or None')

## Training options
flags.DEFINE_integer('pretrain_iterations', 0, 'number of pre-training iterations.')
flags.DEFINE_integer('metatrain_iterations', 15000, 'number of metatraining iterations.') # 60k meta training iterations for miniImagenet (15k for omniglot, 50k for sinusoid)
flags.DEFINE_integer('meta_batch_size', 25, 'number of tasks sampled per meta-update') # 4 for miniImagenet
flags.DEFINE_float('meta_lr', 0.001, 'the base learning rate of the generator')
flags.DEFINE_integer('update_batch_size', 5, 'number of examples used for inner gradient update (K for K-shot learning).')
flags.DEFINE_float('update_lr', 1e-3, 'step size alpha for inner gradient update.') # 0.01 for miniImagenet (0.1 for omniglot)
flags.DEFINE_integer('num_updates', 1, 'number of inner gradient updates during training.') # 5 inner gradient updates for miniImagenet

## Model options
flags.DEFINE_string('model', 'r2d2', 'r2d2 or maml')
flags.DEFINE_string('norm', 'batch_norm', 'batch_norm, layer_norm, or None')
flags.DEFINE_integer('num_filters', 64, 'number of filters for conv nets -- 32 for miniimagenet, 64 for omiglot.') # 32 filters for miniImagenet
flags.DEFINE_bool('conv', True, 'whether or not to use a convolutional network, only applicable in some cases')
flags.DEFINE_bool('max_pool', False, 'Whether or not to use max pooling rather than strided convolutions') # True max pooling for miniImagenet
flags.DEFINE_bool('stop_grad', False, 'if True, do not use second derivatives in meta-optimization (for speed)')

## Logging, saving, and testing options
flags.DEFINE_bool('log', True, 'if false, do not log summaries, for debugging code.')
flags.DEFINE_string('logdir', '/tmp/data', 'directory for summaries and checkpoints.')
flags.DEFINE_bool('resume', True, 'resume training if there is a model available')
flags.DEFINE_bool('train', True, 'True to train, False to test.')
flags.DEFINE_integer('test_iter', -1, 'iteration to load model (-1 for latest model)')
flags.DEFINE_bool('test_set', False, 'Set to true to test on the the test set, False for the validation set.')
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(model, saver, sess, exp_string, data_generator, resume_itr=0):
    """Trains a model with meta learning
        
        Args:
            model:              The class object which is the model we are training on
            saver:              TensorFlow saver object to keep last (10) trainable variables
            sess:               TensorFlow session object signifying the session on which is trained
            exp_string:         String which is used as a folder name to export results to
            data_generator:     data_generator object that generates the right data for the meta learning problem at hand
            resume_itr:         Integer equal to the iteration from which training should be resumed, default = 0
        """
    SUMMARY_INTERVAL = 100
    SAVE_INTERVAL = 1000
    if FLAGS.datasource == 'sinusoid':
        PRINT_INTERVAL = 1000
        TEST_PRINT_INTERVAL = PRINT_INTERVAL*5
    else:
        PRINT_INTERVAL = 100
        TEST_PRINT_INTERVAL = PRINT_INTERVAL*5 # print (1) test eval result only after the train results are printed 5 times

    if FLAGS.log:
        train_writer = tf.summary.FileWriter(FLAGS.logdir + '/' + exp_string, sess.graph)
    print('Done initializing, starting training.')
    prelosses, postlosses = [], []

    num_classes = data_generator.num_classes # for classification, 1 otherwise
    multitask_weights, reg_weights = [], []
    
    meta_lr_damped = FLAGS.meta_lr
    
    # Start iterations from resume_itr if there is a training history
    for itr in range(resume_itr, FLAGS.pretrain_iterations + FLAGS.metatrain_iterations):
        feed_dict = {}
        if 'generate' in dir(data_generator): # This is for sinusoid only
            batch_x, batch_y, amp, phase = data_generator.generate()

            if FLAGS.baseline == 'oracle': # NOTE - this flag is specific to sinusoid
                batch_x = np.concatenate([batch_x, np.zeros([batch_x.shape[0], batch_x.shape[1], 2])], 2)
                for i in range(FLAGS.meta_batch_size):
                    batch_x[i, :, 1] = amp[i]
                    batch_x[i, :, 2] = phase[i]
            
            # a = base-train, b = base-test
            inputa = batch_x[:, :num_classes*FLAGS.update_batch_size, :]
            labela = batch_y[:, :num_classes*FLAGS.update_batch_size, :]
            inputb = batch_x[:, num_classes*FLAGS.update_batch_size:, :]
            labelb = batch_y[:, num_classes*FLAGS.update_batch_size:, :]
            feed_dict = {model.inputa: inputa, model.inputb: inputb,  model.labela: labela, model.labelb: labelb}

        if itr < FLAGS.pretrain_iterations:
            input_tensors = [model.pretrain_op]
        else:
            input_tensors = [model.metatrain_op] # metatrain_op is a tf Operation that does the meta update with Adam

        if (itr % SUMMARY_INTERVAL == 0 or itr % PRINT_INTERVAL == 0):
            # Add all the ops together in one big list, to evaluate them, and print I guess
            input_tensors.extend([model.summ_op, model.total_loss1, model.total_losses2[FLAGS.num_updates-1]])
            if model.classification:
                input_tensors.extend([model.total_accuracy1, model.total_accuracies2[FLAGS.num_updates-1]])
        
        # Do one full meta train step
        result = sess.run(input_tensors, feed_dict)

        if itr % SUMMARY_INTERVAL == 0:
            prelosses.append(result[-2])
            if FLAGS.log:
                train_writer.add_summary(result[1], itr)
            postlosses.append(result[-1])

        if (itr!=0) and itr % PRINT_INTERVAL == 0:
            if itr < FLAGS.pretrain_iterations:
                print_str = 'Pretrain Iteration ' + str(itr)
            else:
                print_str = 'Iteration ' + str(itr - FLAGS.pretrain_iterations)
            print_str += ': ' + str(np.mean(prelosses)) + ', ' + str(np.mean(postlosses))
            print(print_str)
            prelosses, postlosses = [], []

        if (itr!=0) and itr % SAVE_INTERVAL == 0:
            saver.save(sess, FLAGS.logdir + '/' + exp_string + '/model' + str(itr))

        if (itr+1) % 2000 == 0:
            meta_lr_damped = meta_lr_damped*0.5
        
        # sinusoid is infinite data, so no need to test on meta-validation set.
        if (itr!=0) and itr % TEST_PRINT_INTERVAL == 0 and FLAGS.datasource !='sinusoid':
            if 'generate' not in dir(data_generator):
                feed_dict = {}
                if model.classification:
                    input_tensors = [model.metaval_total_accuracy1, model.metaval_total_accuracies2[FLAGS.num_updates-1], model.summ_op]
                else:
                    input_tensors = [model.metaval_total_loss1, model.metaval_total_losses2[FLAGS.num_updates-1], model.summ_op]
                
                feed_dict = {model.meta_lr: meta_lr_damped}
            else:
                batch_x, batch_y, amp, phase = data_generator.generate(train=False)
                inputa = batch_x[:, :num_classes*FLAGS.update_batch_size, :]
                inputb = batch_x[:, num_classes*FLAGS.update_batch_size:, :]
                labela = batch_y[:, :num_classes*FLAGS.update_batch_size, :]
                labelb = batch_y[:, num_classes*FLAGS.update_batch_size:, :]
                feed_dict = {model.inputa: inputa, model.inputb: inputb,  model.labela: labela, model.labelb: labelb, model.meta_lr: 0.0}
                if model.classification:
                    input_tensors = [model.total_accuracy1, model.total_accuracies2[FLAGS.num_updates-1]]
                else:
                    input_tensors = [model.total_loss1, model.total_losses2[FLAGS.num_updates-1]]
            
            # This session run is to evaluate
            result = sess.run(input_tensors, feed_dict)
            print('Validation results: ' + str(result[0]) + ', ' + str(result[1]))

    saver.save(sess, FLAGS.logdir + '/' + exp_string +  '/model' + str(itr))

# calculated for omniglot
NUM_TEST_POINTS = 600

def test(model, saver, sess, exp_string, data_generator, test_num_updates=None):
    """Tests a meta-learned model
        
        Args:
            model:              The class object which is the model we are training on
            saver:              TensorFlow saver object to keep last (10) trainable variables
            sess:               TensorFlow session object signifying the session on which is trained
            exp_string:         String which is used as a folder name to export results to
            data_generator:     data_generator object that generates the right data for the meta learning problem at hand
            test_num_updates:   How many updates are done during testing
        """
    num_classes = data_generator.num_classes # for classification, 1 otherwise

    np.random.seed(1)
    random.seed(1)

    metaval_accuracies = []
    metaval_labels = []

    for _ in range(NUM_TEST_POINTS): # NUM_TEST_POINTS = amount of test tasks
        if 'generate' not in dir(data_generator):
            feed_dict = {}
            feed_dict = {model.meta_lr : 0.0}
        else: # for sinusoid
            batch_x, batch_y, amp, phase = data_generator.generate(train=False)

            if FLAGS.baseline == 'oracle': # NOTE - this flag is specific to sinusoid
                batch_x = np.concatenate([batch_x, np.zeros([batch_x.shape[0], batch_x.shape[1], 2])], 2)
                batch_x[0, :, 1] = amp[0]
                batch_x[0, :, 2] = phase[0]

            inputa = batch_x[:, :num_classes*FLAGS.update_batch_size, :]
            inputb = batch_x[:,num_classes*FLAGS.update_batch_size:, :]
            labela = batch_y[:, :num_classes*FLAGS.update_batch_size, :]
            labelb = batch_y[:,num_classes*FLAGS.update_batch_size:, :]

            feed_dict = {model.inputa: inputa, model.inputb: inputb,  model.labela: labela, model.labelb: labelb, model.meta_lr: 0.0}
        result = sess.run([model.test_accuraciesa] + model.test_accuraciesb, feed_dict)
        labelas, labelbs = sess.run([model.labelas, model.labelbs], feed_dict)
        metaval_accuracies.append(result)
        metaval_labels.append([labelas, labelbs])

    metaval_accuracies = np.array(metaval_accuracies)
    metaval_labels = np.sum(np.array(metaval_labels), axis=0)
    
    means = np.mean(metaval_accuracies, 0)
    stds = np.std(metaval_accuracies, 0)
    ci95 = 1.96*stds/np.sqrt(NUM_TEST_POINTS)
    
    #print(metaval_accuracies)
    print(metaval_labels)
    print('Mean validation accuracy/loss, stddev, and confidence intervals')
    print((means, stds, ci95))

    out_filename = FLAGS.logdir +'/'+ exp_string + '/' + 'test_ubs' + str(FLAGS.update_batch_size) + '_stepsize' + str(FLAGS.update_lr) + '.csv'
    out_pkl = FLAGS.logdir +'/'+ exp_string + '/' + 'test_ubs' + str(FLAGS.update_batch_size) + '_stepsize' + str(FLAGS.update_lr) + '.pkl'
    with open(out_pkl, 'wb') as f:
        pickle.dump({'mses': metaval_accuracies}, f)
    with open(out_filename, 'w') as f:
        writer = csv.writer(f, delimiter=',')
        writer.writerow(['update'+str(i) for i in range(len(means))])
        writer.writerow(means)
        writer.writerow(stds)
        writer.writerow(ci95)

def main():
    """ Puts everything in place to meta-learn and test """
    test_num_updates = 1 # Base learner is linear regression, so only one step required

    if FLAGS.train == False:
        orig_meta_batch_size = FLAGS.meta_batch_size
        # always use meta batch size of 1 when testing.
        FLAGS.meta_batch_size = 1

    if FLAGS.datasource == 'sinusoid':
        # DataGenerator(num_samples_per_class, batch_size, config={})
        data_generator = DataGenerator(FLAGS.update_batch_size*2, FLAGS.meta_batch_size) 
    else: # Dealing with a non 'sinusoid' dataset here
        if FLAGS.metatrain_iterations == 0 and (FLAGS.datasource == 'miniimagenet' or FLAGS.datasource == 'cifarfs'):
            assert FLAGS.meta_batch_size == 1
            assert FLAGS.update_batch_size == 1
            data_generator = DataGenerator(1, FLAGS.meta_batch_size)  # only use one datapoint,
        else:
            if FLAGS.datasource == 'miniimagenet' or FLAGS.datasource == 'cifarfs': # use 15 val examples
                if FLAGS.train: # following Ravi: "15 examples per class were used for evaluating the post-update meta-gradient"
                    # DataGenerator(number_of_images_per_class, number_of_tasks_in_batch)
                    data_generator = DataGenerator(FLAGS.update_batch_size+15, FLAGS.meta_batch_size)
                else: # we're in the testing phase (not train), FLAGS.meta_batch_size = 1
                    data_generator = DataGenerator(FLAGS.update_batch_size*2, FLAGS.meta_batch_size)
            else: # this is for omniglot
                data_generator = DataGenerator(FLAGS.update_batch_size*2, FLAGS.meta_batch_size)


    dim_output = data_generator.dim_output # number of classes, e.g. 5 for miniImagenet tasks
    dim_input = data_generator.dim_input # np.prod(self.img_size) for images

    if FLAGS.datasource == 'miniimagenet' or FLAGS.datasource == 'cifarfs':
        tf_data_load = True
        num_classes = data_generator.num_classes

        if FLAGS.train: # only construct training model if needed
            # meta train : num_total_batches = 200000 (number of tasks, not number of meta-iterations)
            random.seed(5)
            image_tensor, label_tensor = data_generator.make_data_tensor()
            inputa = tf.slice(image_tensor, [0,0,0], [-1,num_classes*FLAGS.update_batch_size, -1]) # slice(tensor, begin, slice_size)
            inputb = tf.slice(image_tensor, [0,num_classes*FLAGS.update_batch_size, 0], [-1,-1,-1]) # The extra 15 add here
            labela = tf.slice(label_tensor, [0,0,0], [-1,num_classes*FLAGS.update_batch_size, -1])
            labelb = tf.slice(label_tensor, [0,num_classes*FLAGS.update_batch_size, 0], [-1,-1,-1])
            input_tensors = {'inputa': inputa, 'inputb': inputb, 'labela': labela, 'labelb': labelb}

        # meta val: num_total_batches = 600 (number of tasks, not number of meta-iterations)
        random.seed(6)
        image_tensor, label_tensor = data_generator.make_data_tensor(train=False)
        inputa = tf.slice(image_tensor, [0,0,0], [-1,num_classes*FLAGS.update_batch_size, -1]) # slice the training examples here
        inputb = tf.slice(image_tensor, [0,num_classes*FLAGS.update_batch_size, 0], [-1,-1,-1]) 
        labela = tf.slice(label_tensor, [0,0,0], [-1,num_classes*FLAGS.update_batch_size, -1])
        labelb = tf.slice(label_tensor, [0,num_classes*FLAGS.update_batch_size, 0], [-1,-1,-1])
        metaval_input_tensors = {'inputa': inputa, 'inputb': inputb, 'labela': labela, 'labelb': labelb}
    else:
        tf_data_load = False
        input_tensors = None

    model = R2D2_paper(dim_input, dim_output, test_num_updates=test_num_updates) # test_num_updates = eval on at least one update for training, 10 testing
    if FLAGS.train or not tf_data_load:
        model.construct_model(input_tensors=input_tensors, prefix='metatrain_')
    if tf_data_load:
        model.construct_model(input_tensors=metaval_input_tensors, prefix='metaval_')
    
    # Op to retrieve summaries
    model.summ_op = tf.summary.merge_all()
    
    # keep last 10 copies of trainable variables
    saver = loader = tf.train.Saver(tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES), max_to_keep=10)
    
    # remove the need to explicitly pass this Session object to run ops
    sess = tf.InteractiveSession()

    if FLAGS.train == False:
        # change to original meta batch size when loading model.
        FLAGS.meta_batch_size = orig_meta_batch_size

    if FLAGS.train_update_batch_size == -1:
        FLAGS.train_update_batch_size = FLAGS.update_batch_size
    if FLAGS.train_update_lr == -1:
        FLAGS.train_update_lr = FLAGS.update_lr
    
    # cls = no of classes
    # mbs = meta batch size
    # ubs = update batch size
    # numstep = number of INNER GRADIENT updates
    # updatelr = inner gradient step
    exp_string = 'cls_'+str(FLAGS.num_classes)+'.mbs_'+str(FLAGS.meta_batch_size) + '.ubs_' + str(FLAGS.train_update_batch_size) + '.numstep' + str(FLAGS.num_updates) + '.updatelr' + str(FLAGS.train_update_lr)

    if FLAGS.num_filters != 64:
        exp_string += 'hidden' + str(FLAGS.num_filters)
    if FLAGS.max_pool:
        exp_string += 'maxpool'
    if FLAGS.stop_grad:
        exp_string += 'stopgrad'
    if FLAGS.baseline:
        exp_string += FLAGS.baseline
    if FLAGS.norm == 'batch_norm':
        exp_string += 'batchnorm'
    elif FLAGS.norm == 'layer_norm':
        exp_string += 'layernorm'
    elif FLAGS.norm == 'None':
        exp_string += 'nonorm'
    else:
        print('Norm setting not recognized.')

    resume_itr = 0
    model_file = None
    
    # Initialize all variables
    tf.global_variables_initializer().run()
    # starts threads for all queue runners collected in the graph
    tf.train.start_queue_runners()

    if FLAGS.resume or not FLAGS.train:
        model_file = tf.train.latest_checkpoint(FLAGS.logdir + '/' + exp_string)
        if FLAGS.test_iter > 0:
            model_file = model_file[:model_file.index('model')] + 'model' + str(FLAGS.test_iter)
        if model_file:
            ind1 = model_file.index('model')
            resume_itr = int(model_file[ind1+5:])
            print("Restoring model weights from " + model_file)
            saver.restore(sess, model_file)

    if FLAGS.train:
        train(model, saver, sess, exp_string, data_generator, resume_itr)
    else:
        test(model, saver, sess, exp_string, data_generator, test_num_updates)

if __name__ == "__main__":
    main()