train.py

import os
import sys
from collections import defaultdict
import pickle
import json
import time

import numpy as np
import tensorflow as tf

import task
from model import SingleLayerModel, FullModel, NormalizedMLP, AutoEncoder, AutoEncoderSimple, RNN, ParameterizeK
from configs import FullConfig, SingleLayerConfig
import tools
from standard.analysis_pn2kc_training import _compute_sparsity


def make_input(x, y, batch_size):
    data = tf.data.Dataset.from_tensor_slices((x, y))
    data = data.shuffle(int(1E6))
    # Making sure the shape is fully defined
    try:
        data = data.batch(tf.cast(batch_size, tf.int64), drop_remainder=True)
    except TypeError:
        data = data.apply(tf.contrib.data.batch_and_drop_remainder(batch_size))
    # data = data.batch(tf.cast(batch_size, tf.int64))
    data = data.repeat()
    train_iter = data.make_initializable_iterator()
    next_element = train_iter.get_next()
    return train_iter, next_element


def train_from_path(path):
    """Train from a path with a config file in it."""
    config = tools.load_config(path)
    train(config, reload=True)


def train(config, reload=False, save_everytrainloss=False):
    tf.reset_default_graph()

    # Merge model config with config from dataset
    dataset_config = tools.load_config(config.data_dir)
    dataset_config.update(config)
    config = dataset_config
    for item in config.__dict__.items():
        print(item)

    if not os.path.exists(config.save_path):
        os.makedirs(config.save_path)
    # Save config
    tools.save_config(config, save_path=config.save_path)

    # Load dataset
    train_x, train_y, val_x, val_y = task.load_data(config.data_dir)

    batch_size = config.batch_size
    if 'n_batch' in dir(config):
        n_batch = config.n_batch
    else:
        n_batch = train_x.shape[0] // batch_size

    if config.model == 'full':
        CurrentModel = FullModel
    elif config.model == 'singlelayer':
        CurrentModel = SingleLayerModel
    elif config.model == 'normmlp':
        CurrentModel = NormalizedMLP
    elif config.model == 'oracle':
        CurrentModel = OracleNet
    elif config.model == 'autoencode':
        CurrentModel = AutoEncoder
        # CurrentModel = AutoEncoderSimple
    elif config.model == 'rnn':
        CurrentModel = RNN
    elif config.model == 'K':
        CurrentModel = ParameterizeK
    else:
        raise ValueError('Unknown model type ' + str(config.model))

    # Build train model
    train_x_ph = tf.placeholder(train_x.dtype, train_x.shape)
    train_y_ph = tf.placeholder(train_y.dtype, train_y.shape)
    train_iter, next_element = make_input(train_x_ph, train_y_ph, batch_size)
    model = CurrentModel(next_element[0], next_element[1], config=config)

    # Build validation model
    val_x_ph = tf.placeholder(val_x.dtype, val_x.shape)
    val_y_ph = tf.placeholder(val_y.dtype, val_y.shape)
    val_model = CurrentModel(val_x_ph, val_y_ph, config=config, training=False)

    if 'set_oracle' in dir(config) and config.set_oracle:
        # Helper model for oracle
        oracle_x_ph = tf.placeholder(val_x.dtype, [config.N_CLASS, val_x.shape[1]])
        oracle_y_ph = tf.placeholder(val_y.dtype, [config.N_CLASS])
        oracle = CurrentModel(oracle_x_ph, oracle_y_ph, config=config, training=False)

    # Make custom logger
    log = defaultdict(list)

    glo_score_mode = 'tile' if config.replicate_orn_with_tiling else 'repeat'

    # validation fetches
    val_fetch_names = ['loss', 'acc']
    try:
        _ = val_model.acc2
        val_fetch_names.append('acc2')
    except AttributeError:
        pass

    if config.model == 'full' and config.train_pn2kc:
        val_fetch_names.append('kc')

    val_fetches = [getattr(val_model, f) for f in val_fetch_names]

    tf_config = tf.ConfigProto()
    tf_config.gpu_options.allow_growth = True
    tf_config.log_device_placement = False

    finish_training = False
    with tf.Session(config=tf_config) as sess:
        # sess.run(tf.global_variables_initializer())
        sess.run(tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()))
        sess.run(train_iter.initializer, feed_dict={train_x_ph: train_x,
                                                    train_y_ph: train_y})


        start_epoch = 0
        if reload:
            try:
                model.load()
                log = tools.load_log(config.save_path)
                start_epoch = log['epoch'][-1] + 1
            except:
                print('No model file to be reloaded, starting anew')

        if 'set_oracle' in dir(config) and config.set_oracle:
            oracle.set_oracle_weights()

        if config.model == 'rnn' and config.DIAGONAL:
            model.set_weights()

        loss = 0
        acc = 0
        lr = config.lr
        acc_smooth = 0
        total_time, start_time = 0, time.time()
        w_bins = np.linspace(0, 1, 201)
        w_bins_log = np.linspace(-20, 5, 201)
        log['w_bins'] = w_bins
        log['w_bins_log'] = w_bins_log
        lin_bins = np.linspace(0, 1, 1001)
        log['lin_bins'] = lin_bins
        activity_bins = np.linspace(0, 1, 201)

        for ep in range(start_epoch, config.max_epoch):
            # Validation
            tmp = sess.run(val_fetches, {val_x_ph: val_x, val_y_ph: val_y})
            res = {name:r for name, r in zip(val_fetch_names, tmp)}

            if ep % config.save_epoch_interval == 0:
                print('[*' + '*'*50 + '*]')
                print('Epoch {:d}'.format(ep))
                print('Train/Validation loss {:0.2f}/{:0.2f}'.format(loss, res['loss']))
                print('Train/Validation accuracy {:0.2f}/{:0.2f}'.format(acc, res['acc']))
                print('Learning rate {:0.3E}'.format(lr))
                log['epoch'].append(ep)
                log['train_loss'].append(loss)
                log['train_acc'].append(acc)
                log['lr_now'].append(lr)
                for key, value in res.items():
                    if key in ['kc']:
                        continue
                    log['val_' + key].append(value)

                try:
                    print('Validation accuracy head 2 {:0.2f}'.format(res['acc2']))
                except KeyError:
                    pass

                if config.model == 'full':
                    if config.train_pn2kc:
                        w_glo = sess.run(model.w_glo)
                        w_glo[w_glo<1e-9] = 1e-9 #finite range for log
                        kcs = res['kc']

                        coding_level = (kcs > 0).mean()
                        coding_level_per_kc = kcs.mean(axis=0)
                        coding_level_per_odor = kcs.mean(axis=1)
                        log['coding_level'].append(coding_level)
                        hist, _ = np.histogram(coding_level_per_kc, bins=activity_bins)
                        log['coding_level_per_kc'].append(hist)
                        hist, _ = np.histogram(coding_level_per_odor, bins=activity_bins)
                        log['coding_level_per_odor'].append(hist)

                        # Store distribution of flattened weigths
                        log_hist, _ = np.histogram(np.log(w_glo.flatten()), bins=w_bins_log)
                        hist, _ = np.histogram(w_glo.flatten(), bins=w_bins)
                        log['log_hist'].append(log_hist)
                        log['hist'].append(hist)
                        log['kc_w_sum'].append(w_glo.sum(axis=0))
                        hist, _ = np.histogram(w_glo.flatten(), bins=lin_bins)
                        log['lin_hist'].append(hist)
                        # Store sparsity computed with threshold

                        sparsity_inferred, thres_inferred = _compute_sparsity(w_glo, dynamic_thres=True, thres=.1)
                        K_inferred = sparsity_inferred[sparsity_inferred > 0 ].mean()
                        bad_KC_inferred = np.sum(sparsity_inferred == 0) / sparsity_inferred.size
                        log['sparsity_inferred'].append(sparsity_inferred)
                        log['thres_inferred'].append(thres_inferred)
                        log['K_inferred'].append(K_inferred)
                        log['bad_KC_inferred'].append(bad_KC_inferred)

                        sparsity, thres = _compute_sparsity(w_glo, dynamic_thres=False, thres= config.kc_prune_threshold)
                        K = sparsity[sparsity > 0].mean()
                        bad_KC = np.sum(sparsity == 0)/sparsity.size
                        log['sparsity'].append(sparsity)
                        log['thres'].append(thres)
                        log['K'].append(K)
                        log['bad_KC'].append(bad_KC)

                        print('KC coding level={}'.format(np.round(coding_level,2)))
                        print('Bad KCs (fixed, inferred) ={}, {}'.format(bad_KC, bad_KC_inferred))
                        print('K (fixed, inferred) ={}, {}'.format(K, K_inferred))

                    if config.receptor_layer:
                        w_or = sess.run(model.w_or)
                        or_glo_score, _ = tools.compute_glo_score(
                            w_or, config.N_ORN, glo_score_mode)
                        print('OR receptor glo score ' + str(or_glo_score))
                        log['or_glo_score'].append(or_glo_score)

                        w_orn = sess.run(model.w_orn)
                        glo_score, _ = tools.compute_glo_score(
                            w_orn, config.N_ORN, 'matrix', w_or)
                        print('Glo score ' + str(glo_score))
                        log['glo_score'].append(glo_score)

                        w_combined = np.matmul(w_or, w_orn)
                        combined_glo_score, _ = tools.compute_glo_score(
                            w_combined, config.N_ORN, glo_score_mode)
                        print('Combined glo score ' + str(combined_glo_score))
                        log['combined_glo_score'].append(combined_glo_score)

                    else:
                        if config.train_orn2pn and not config.skip_orn2pn:
                            w_orn = sess.run(model.w_orn)
                            glo_score, _ = tools.compute_glo_score(
                                w_orn, config.N_ORN, glo_score_mode)
                            log['glo_score'].append(glo_score)
                            print('Glo score ' + str(glo_score))

                            sim_score, _ = tools.compute_sim_score(
                                w_orn, config.N_ORN, glo_score_mode)
                            log['sim_score'].append(sim_score)
                            print('Sim score ' + str(sim_score))
                        else:
                            w_orn = sess.run(model.w_orn)
                            glo_score, _ = tools.compute_glo_score(
                                w_orn, config.N_ORN, glo_score_mode)
                            log['glo_score'].append(glo_score)
                            print('Glo score ' + str(glo_score))

                elif config.model == 'K':
                    K = sess.run(model.K)
                    bias = sess.run(model.b_glo)
                    kcs = sess.run(model.kc)

                    log['K'].append(K)
                    print('K ={}'.format(np.round(K,2)))
                    print('Bias (mean)={}'.format(np.round(bias.mean(),2)))
                    print('kc activity (mean) = {}'.format(kcs.mean()))
                    print('kc sparseness (mean) = {}'.format(np.mean(kcs > 0)))
                    # import matplotlib.pyplot as plt
                    # plt.imshow(w_glo[:,:30], vmin=0, vmax=2)
                    # plt.colorbar()
                    # plt.show()

                        # w_glo = sess.run(model.w_glo)
                        # glo_score_w_glo, _ = tools.compute_glo_score(w_glo)
                        # log['glo_score_w_glo'].append(glo_score_w_glo, config.N_ORN)

                # Compute condition number
                # w_glo = sess.run(model.w_glo)
                # w_orn2kc = np.dot(w_orn, w_glo)
                # cond = np.linalg.cond(w_orn2kc)
                # log['cond'].append(cond)
                # print('Condition number '+ str(cond))

                if ep > 0:
                    time_spent = time.time() - start_time
                    total_time += time_spent
                    print('Time taken {:0.1f}s'.format(total_time))
                    print('Examples/second {:d}'.format(int(train_x.shape[0]/time_spent)))
                start_time = time.time()

                tools.save_log(config.save_path, log)

            if 'target_acc' in dir(config) and config.target_acc is not None:
                if res['acc'] > config.target_acc:
                    print('Training reached target accuracy {:0.2f}>{:0.2f}'.format(
                        res['acc'], config.target_acc
                    ))
                    finish_training = True

            try:
                if config.save_every_epoch and ep % config.save_epoch_interval == 0:
                    model.save_pickle(ep)
                    model.save(ep)

                # Train
                if save_everytrainloss:
                    for b in range(n_batch-1):
                        if config.separate_optimizer:
                            loss, acc, _, _ = sess.run([model.loss, model.acc, model.train_op, model.train_op1])
                        else:
                            loss, acc, _ = sess.run([model.loss, model.acc, model.train_op])
                        log['train_loss'].append(loss)

                        acc_smooth = acc_smooth * 0.9 + acc * 0.1
                        if config.target_acc is not None and acc_smooth > config.target_acc:
                            print(
                                'Training reached target accuracy {:0.2f}>{:0.2f}'.format(
                                    acc_smooth, config.target_acc
                                ))
                            finish_training = True
                            break
                else:
                    for b in range(n_batch-1):
                        _ = sess.run(model.train_op)

                        # if b % 10 == 0:
                            # w_orn, w_glo = sess.run([model.w_orn, model.w_glo])
                            # weights_over_time.append((w_orn, w_glo))

                # Compute training loss and accuracy using last batch
                loss, acc, _ = sess.run([model.loss, model.acc, model.train_op])

            except KeyboardInterrupt:
                print('Training interrupted by users')
                finish_training = True

            if finish_training:
                break

            sys.stdout.flush()

        print('Training finished')
        if 'save_log_only' in dir(config) and config.save_log_only is True:
            pass
        else:
            model.save_pickle()
            model.save()



if __name__ == '__main__':
    experiment = 'robert'
    # experiment = 'peter'
    if experiment == 'peter':
        config = SingleLayerConfig()

    elif experiment == 'robert':
        config = FullConfig()
        config.dataset = 'proto'
        config.data_dir = './datasets/proto/standard'
        config.model = 'full'
        config.save_path = './files/peter'
        config.separate_optimizer = True
        config.train_pn2kc = True
        config.sparse_pn2kc = False
    else:
        raise NotImplementedError

    train(config)