--- a
+++ b/EEGLearn/train.py
@@ -0,0 +1,355 @@
+##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+## Created by: Yang Wang
+## School of Automation, Huazhong University of Science & Technology (HUST)
+## wangyang_sky@hust.edu.cn
+## Copyright (c) 2018
+##
+## This source code is licensed under the MIT-style license found in the
+## LICENSE file in the root directory of this source tree
+##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+
+#coding:utf-8
+
+import os
+import tensorflow as tf
+import numpy as np
+import scipy.io
+import time
+import datetime
+
+from utils import reformatInput, load_or_generate_images, iterate_minibatches
+
+from model import build_cnn, build_convpool_conv1d, build_convpool_lstm, build_convpool_mix
+
+
+timestamp = datetime.datetime.now().strftime('%Y-%m-%d.%H.%M')
+log_path = os.path.join("runs", timestamp)
+
+
+model_type = '1dconv'      # ['1dconv', 'maxpool', 'lstm', 'mix', 'cnn']
+log_path = log_path + '_' + model_type
+
+batch_size = 32
+dropout_rate = 0.5
+
+input_shape = [32, 32, 3]   # 1024
+nb_class = 4
+n_colors = 3
+
+# whether to train cnn first, and load its weight for multi-frame model
+reuse_cnn_flag = False
+
+# learning_rate for different models
+lrs = {
+    'cnn': 1e-3,
+    '1dconv': 1e-4,
+    'lstm': 1e-4,
+    'mix': 1e-4,
+}
+
+weight_decay = 1e-4
+learning_rate = lrs[model_type] / 32 * batch_size
+optimizer = tf.train.AdamOptimizer
+
+num_epochs = 60
+
+def train(images, labels, fold, model_type, batch_size, num_epochs, subj_id=0, reuse_cnn=False, 
+    dropout_rate=dropout_rate ,learning_rate_default=1e-3, Optimizer=tf.train.AdamOptimizer, log_path=log_path):
+    """
+    A sample training function which loops over the training set and evaluates the network
+    on the validation set after each epoch. Evaluates the network on the training set
+    whenever the
+    :param images: input images
+    :param labels: target labels
+    :param fold: tuple of (train, test) index numbers
+    :param model_type: model type ('cnn', '1dconv', 'lstm', 'mix')
+    :param batch_size: batch size for training
+    :param num_epochs: number of epochs of dataset to go over for training
+    :param subj_id: the id of fold for storing log and the best model
+    :param reuse_cnn: whether to train cnn first, and load its weight for multi-frame model
+    :return: none
+    """
+
+    with tf.name_scope('Inputs'):
+        input_var = tf.placeholder(tf.float32, [None, None, 32, 32, n_colors], name='X_inputs')
+        target_var = tf.placeholder(tf.int64, [None], name='y_inputs')
+        tf_is_training = tf.placeholder(tf.bool, None, name='is_training')
+
+    num_classes = len(np.unique(labels))
+    (X_train, y_train), (X_val, y_val), (X_test, y_test) = reformatInput(images, labels, fold)
+
+
+    print('Train set label and proportion:\t', np.unique(y_train, return_counts=True))
+    print('Val   set label and proportion:\t', np.unique(y_val, return_counts=True))
+    print('Test  set label and proportion:\t', np.unique(y_test, return_counts=True))
+
+    print('The shape of X_trian:\t', X_train.shape)
+    print('The shape of X_val:\t', X_val.shape)
+    print('The shape of X_test:\t', X_test.shape)
+    
+
+    print("Building model and compiling functions...")
+    if model_type == '1dconv':
+        network = build_convpool_conv1d(input_var, num_classes, train=tf_is_training, 
+                            dropout_rate=dropout_rate, name='CNN_Conv1d'+'_sbj'+str(subj_id))
+    elif model_type == 'lstm':
+        network = build_convpool_lstm(input_var, num_classes, 100, train=tf_is_training, 
+                            dropout_rate=dropout_rate, name='CNN_LSTM'+'_sbj'+str(subj_id))
+    elif model_type == 'mix':
+        network = build_convpool_mix(input_var, num_classes, 100, train=tf_is_training, 
+                            dropout_rate=dropout_rate, name='CNN_Mix'+'_sbj'+str(subj_id))
+    elif model_type == 'cnn':
+        with tf.name_scope(name='CNN_layer'+'_fold'+str(subj_id)):
+            network = build_cnn(input_var)  # output shape [None, 4, 4, 128]
+            convpool_flat = tf.reshape(network, [-1, 4*4*128])
+            h_fc1_drop1 = tf.layers.dropout(convpool_flat, rate=dropout_rate, training=tf_is_training, name='dropout_1')
+            h_fc1 = tf.layers.dense(h_fc1_drop1, 256, activation=tf.nn.relu, name='fc_relu_256')
+            h_fc1_drop2 = tf.layers.dropout(h_fc1, rate=dropout_rate, training=tf_is_training, name='dropout_2')
+            network = tf.layers.dense(h_fc1_drop2, num_classes, name='fc_softmax')
+            # the loss function contains the softmax activation
+    else:
+        raise ValueError("Model not supported ['1dconv', 'maxpool', 'lstm', 'mix', 'cnn']")
+
+    Train_vars = tf.trainable_variables()
+
+    prediction = network
+
+    with tf.name_scope('Loss'):
+        l2_loss = tf.add_n([tf.nn.l2_loss(v) for v in Train_vars if 'kernel' in v.name])
+        ce_loss = tf.losses.sparse_softmax_cross_entropy(labels=target_var, logits=prediction)
+        _loss = ce_loss + weight_decay*l2_loss
+
+    # decay_steps learning rate decay
+    decay_steps = 3*(len(y_train)//batch_size)   # len(X_train)//batch_size  the training steps for an epcoh
+    with tf.name_scope('Optimizer'):
+        # learning_rate = learning_rate_default * Decay_rate^(global_steps/decay_steps)
+        global_steps = tf.Variable(0, name="global_step", trainable=False)
+        learning_rate = tf.train.exponential_decay(     # learning rate decay
+            learning_rate_default,  # Base learning rate.
+            global_steps,
+            decay_steps,
+            0.95,  # Decay rate.
+            staircase=True)
+        optimizer = Optimizer(learning_rate)    # GradientDescentOptimizer  AdamOptimizer
+        train_op = optimizer.minimize(_loss, global_step=global_steps, var_list=Train_vars)
+
+    with tf.name_scope('Accuracy'):
+        prediction = tf.argmax(prediction, axis=1)
+        correct_prediction = tf.equal(prediction, target_var)
+        accuracy = tf.reduce_mean(tf.cast(correct_prediction,tf.float32))
+
+    # Output directory for models and summaries
+    # choose different path for different model and subject
+    out_dir = os.path.abspath(os.path.join(os.path.curdir, log_path, (model_type+'_'+str(subj_id)) ))
+    print("Writing to {}\n".format(out_dir))
+
+    # Summaries for loss, accuracy and learning_rate
+    loss_summary = tf.summary.scalar('loss', _loss)
+    acc_summary = tf.summary.scalar('train_acc', accuracy)
+    lr_summary = tf.summary.scalar('learning_rate', learning_rate)
+
+    # Train Summaries
+    train_summary_op = tf.summary.merge([loss_summary, acc_summary, lr_summary])
+    train_summary_dir = os.path.join(out_dir, "summaries", "train")
+    train_summary_writer = tf.summary.FileWriter(train_summary_dir, tf.get_default_graph())
+
+    # Dev summaries
+    dev_summary_op = tf.summary.merge([loss_summary, acc_summary])
+    dev_summary_dir = os.path.join(out_dir, "summaries", "dev")
+    dev_summary_writer = tf.summary.FileWriter(dev_summary_dir, tf.get_default_graph())
+
+    # Test summaries
+    test_summary_op = tf.summary.merge([loss_summary, acc_summary])
+    test_summary_dir = os.path.join(out_dir, "summaries", "test")
+    test_summary_writer = tf.summary.FileWriter(test_summary_dir, tf.get_default_graph())
+
+
+    # Checkpoint directory. Tensorflow assumes this directory already exists so we need to create it
+    checkpoint_dir = os.path.abspath(os.path.join(out_dir, "checkpoints"))
+    checkpoint_prefix = os.path.join(checkpoint_dir, model_type)
+    if not os.path.exists(checkpoint_dir):
+        os.makedirs(checkpoint_dir)
+
+
+    if model_type != 'cnn' and reuse_cnn:
+        # saver for reuse the CNN weight
+        reuse_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='VGG_NET_CNN')
+        original_saver = tf.train.Saver(reuse_vars)         # Pass the variables as a list
+
+    saver = tf.train.Saver(tf.global_variables(), max_to_keep=1)
+
+    print("Starting training...")
+    total_start_time = time.time()
+    best_validation_accu = 0
+
+    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
+    with tf.Session() as sess:
+        sess.run(init_op)
+        if model_type != 'cnn' and reuse_cnn:
+            cnn_model_path = os.path.abspath(
+                                os.path.join(
+                                    os.path.curdir, log_path, ('cnn_'+str(subj_id)), 'checkpoints' ))
+            cnn_model_path = tf.train.latest_checkpoint(cnn_model_path)
+            print('-'*20)
+            print('Load cnn model weight for multi-frame model from {}'.format(cnn_model_path))
+            original_saver.restore(sess, cnn_model_path)
+
+        stop_count = 0  # count for earlystopping
+        for epoch in range(num_epochs):
+            print('-'*50)
+            # Train set
+            train_err = train_acc = train_batches = 0
+            start_time = time.time()
+            for batch in iterate_minibatches(X_train, y_train, batch_size, shuffle=False):
+                inputs, targets = batch
+                summary, _, pred, loss, acc = sess.run([train_summary_op, train_op, prediction, _loss, accuracy], 
+                    {input_var: inputs, target_var: targets, tf_is_training: True})
+                train_acc += acc
+                train_err += loss
+                train_batches += 1
+                train_summary_writer.add_summary(summary, sess.run(global_steps))
+
+            av_train_err = train_err / train_batches
+            av_train_acc = train_acc / train_batches
+
+            # Val set
+            summary, pred, av_val_err, av_val_acc = sess.run([dev_summary_op, prediction, _loss, accuracy],
+                    {input_var: X_val, target_var: y_val, tf_is_training: False})
+            dev_summary_writer.add_summary(summary, sess.run(global_steps))
+
+            
+            print("Epoch {} of {} took {:.3f}s".format(
+                epoch + 1, num_epochs, time.time() - start_time))
+            
+            fmt_str = "Train \tEpoch [{:d}/{:d}]  train_Loss: {:.4f}\ttrain_Acc: {:.2f}"
+            print_str = fmt_str.format(epoch + 1, num_epochs, av_train_err, av_train_acc*100)
+            print(print_str)
+
+            fmt_str = "Val \tEpoch [{:d}/{:d}]  val_Loss: {:.4f}\tval_Acc: {:.2f}"
+            print_str = fmt_str.format(epoch + 1, num_epochs, av_val_err, av_val_acc*100)
+            print(print_str)
+            
+            # Test set
+            summary, pred, av_test_err, av_test_acc = sess.run([test_summary_op, prediction, _loss, accuracy],
+                {input_var: X_test, target_var: y_test, tf_is_training: False})
+            test_summary_writer.add_summary(summary, sess.run(global_steps))
+            
+            fmt_str = "Test \tEpoch [{:d}/{:d}]  test_Loss: {:.4f}\ttest_Acc: {:.2f}"
+            print_str = fmt_str.format(epoch + 1, num_epochs, av_test_err, av_test_acc*100)
+            print(print_str)
+
+            if av_val_acc > best_validation_accu:   # early_stoping
+                stop_count = 0
+                eraly_stoping_epoch = epoch
+                best_validation_accu = av_val_acc
+                test_acc_val = av_test_acc
+                saver.save(sess, checkpoint_prefix, global_step=sess.run(global_steps))
+            else:
+                stop_count += 1
+                if stop_count >= 10: # stop training if val_acc dose not imporve for over 10 epochs
+                    break
+
+        train_batches = train_acc = 0
+        for batch in iterate_minibatches(X_train, y_train, batch_size, shuffle=False):
+            inputs, targets = batch
+            acc = sess.run(accuracy, {input_var: X_train, target_var: y_train, tf_is_training: False})
+            train_acc += acc
+            train_batches += 1
+
+        last_train_acc = train_acc / train_batches
+        
+        
+        last_val_acc = av_val_acc
+        last_test_acc = av_test_acc
+        print('-'*50)
+        print('Time in total:', time.time()-total_start_time)
+        print("Best validation accuracy:\t\t{:.2f} %".format(best_validation_accu * 100))
+        print("Test accuracy when got the best validation accuracy:\t\t{:.2f} %".format(test_acc_val * 100))
+        print('-'*50)
+        print("Last train accuracy:\t\t{:.2f} %".format(last_train_acc * 100))
+        print("Last validation accuracy:\t\t{:.2f} %".format(last_val_acc * 100))
+        print("Last test accuracy:\t\t\t\t{:.2f} %".format(last_test_acc * 100))
+        print('Early Stopping at epoch: {}'.format(eraly_stoping_epoch+1))
+
+    train_summary_writer.close()
+    dev_summary_writer.close()
+    test_summary_writer.close()
+    return [last_train_acc, best_validation_accu, test_acc_val, last_val_acc, last_test_acc]
+
+
+
+def train_all_model(num_epochs=3000):
+    nums_subject = 13
+    # Leave-Subject-Out cross validation
+    subj_nums = np.squeeze(scipy.io.loadmat('../SampleData/trials_subNums.mat')['subjectNum'])
+    fold_pairs = []
+    for i in np.unique(subj_nums):
+        ts = subj_nums == i
+        tr = np.squeeze(np.nonzero(np.bitwise_not(ts)))
+        ts = np.squeeze(np.nonzero(ts))
+        np.random.shuffle(tr)
+        np.random.shuffle(ts)
+        fold_pairs.append((tr, ts))
+
+
+    images_average, images_timewin, labels = load_or_generate_images(
+                                                file_path='../SampleData/', average_image=3)
+
+
+    print('*'*200)
+    acc_buf = []
+    for subj_id in range(nums_subject):
+        print('-'*100)
+        
+        if model_type == 'cnn':
+            print('The subjects', subj_id, '\t\t Training the ' + 'cnn' + ' Model...')
+            acc_temp = train(images_average, labels, fold_pairs[subj_id], 'cnn', 
+                                batch_size=batch_size, num_epochs=num_epochs, subj_id=subj_id,
+                                learning_rate_default=lrs['cnn'], Optimizer=optimizer, log_path=log_path)
+            acc_buf.append(acc_temp)
+            tf.reset_default_graph()
+            print('Done!')
+
+        else:
+            # whether to train cnn first, and load its weight for multi-frame model
+            if reuse_cnn_flag is True:
+                print('The subjects', subj_id, '\t\t Training the ' + 'cnn' + ' Model...')
+                acc_temp = train(images_average, labels, fold_pairs[subj_id], 'cnn', 
+                                    batch_size=batch_size, num_epochs=num_epochs, subj_id=subj_id,
+                                    learning_rate_default=lrs['cnn'], Optimizer=optimizer, log_path=log_path)
+                # acc_buf.append(acc_temp)
+                tf.reset_default_graph()
+                print('Done!')
+        
+            print('The subjects', subj_id, '\t\t Training the ' + model_type + ' Model...')
+            print('Load the CNN model weight for backbone...')
+            acc_temp = train(images_timewin, labels, fold_pairs[subj_id], model_type, 
+                            batch_size=batch_size, num_epochs=num_epochs, subj_id=subj_id, reuse_cnn=reuse_cnn_flag, 
+                            learning_rate_default=learning_rate, Optimizer=optimizer, log_path=log_path)
+                                
+            acc_buf.append(acc_temp)
+            tf.reset_default_graph()
+            print('Done!')
+        
+        # return
+
+    print('All folds for {} are done!'.format(model_type))
+    acc_buf = (np.array(acc_buf)).T
+    acc_mean = np.mean(acc_buf, axis=1).reshape(-1, 1)
+    acc_buf = np.concatenate([acc_buf, acc_mean], axis=1)
+    # the last column is the mean of current row
+    print('Last_train_acc:\t', acc_buf[0], '\tmean :', np.mean(acc_buf[0][-1]))
+    print('Best_val_acc:\t', acc_buf[1], '\tmean :', np.mean(acc_buf[1][-1]))
+    print('Earlystopping_test_acc:\t', acc_buf[2], '\tmean :', np.mean(acc_buf[2][-1]))
+    print('Last_val_acc:\t', acc_buf[3], '\tmean :', np.mean(acc_buf[3][-1]))
+    print('Last_test_acc:\t', acc_buf[4], '\tmean :', np.mean(acc_buf[4][-1]))
+    np.savetxt('./Accuracy_{}.csv'.format(model_type), acc_buf, fmt='%.4f', delimiter=',')
+
+
+if __name__ == '__main__':
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+    np.random.seed(2018)
+    tf.set_random_seed(2018)
+
+    train_all_model(num_epochs=num_epochs)