train_luna_prop_patch.py

import cPickle as pickle
import string
import sys
import time
from itertools import izip
import lasagne as nn
import numpy as np
import theano
from datetime import datetime, timedelta
from collections import defaultdict
import utils
import logger
import theano.tensor as T
import buffering
from configuration import config, set_configuration
import pathfinder

theano.config.warn_float64 = 'raise'

if len(sys.argv) < 2:
    sys.exit("Usage: train.py <configuration_name>")

config_name = sys.argv[1]
set_configuration('configs_luna_props_patch', config_name)
expid = utils.generate_expid(config_name)
print
print "Experiment ID: %s" % expid
print

# metadata
metadata_dir = utils.get_dir_path('models', pathfinder.METADATA_PATH)
metadata_path = metadata_dir + '/%s.pkl' % expid

# logs
logs_dir = utils.get_dir_path('logs', pathfinder.METADATA_PATH)
sys.stdout = logger.Logger(logs_dir + '/%s.log' % expid)
sys.stderr = sys.stdout

print 'Build model'
model = config().build_model()
all_layers = nn.layers.get_all_layers(model.l_out)
all_params = nn.layers.get_all_params(model.l_out)
num_params = nn.layers.count_params(model.l_out)
print '  number of parameters: %d' % num_params
print string.ljust('  layer output shapes:', 36),
print string.ljust('#params:', 10),
print 'output shape:'
for layer in all_layers:
    name = string.ljust(layer.__class__.__name__, 32)
    num_param = sum([np.prod(p.get_value().shape) for p in layer.get_params()])
    num_param = string.ljust(num_param.__str__(), 10)
    print '    %s %s %s %s' % (name, num_param, layer.output_shape, layer.name)

train_loss = config().build_objective(model, deterministic=False)
valid_loss = config().build_objective(model, deterministic=True)

learning_rate_schedule = config().learning_rate_schedule
learning_rate = theano.shared(np.float32(learning_rate_schedule[0]))
updates = config().build_updates(train_loss, model, learning_rate)

x_shared = nn.utils.shared_empty(dim=len(model.l_in.shape))
y_shared = nn.utils.shared_empty(dim=len(model.l_target.shape))
if config().need_enable:
    z_shared = nn.utils.shared_empty(dim=len(model.l_enable_target.shape))

idx = T.lscalar('idx')
givens_train = {}
givens_train[model.l_in.input_var] = x_shared[idx * config().batch_size:(idx + 1) * config().batch_size]
givens_train[model.l_target.input_var] = y_shared[idx * config().batch_size:(idx + 1) * config().batch_size]
if config().need_enable:
    givens_train[model.l_enable_target.input_var] =  z_shared[idx * config().batch_size:(idx + 1) * config().batch_size]

givens_valid = {}
givens_valid[model.l_in.input_var] = x_shared
givens_valid[model.l_target.input_var] = y_shared
# at this moment we do not use the enable target
if config().need_enable:
    givens_valid[model.l_enable_target.input_var] = z_shared


#first make ordered list of objective functions
train_objectives = [config().d_objectives[obj_name] for obj_name in config().order_objectives]
test_objectives = [config().d_objectives_deterministic[obj_name] for obj_name in config().order_objectives]
# theano functions
print givens_train
iter_train = theano.function([idx], train_objectives, givens=givens_train, updates=updates)

print 'test_objectives'
print config().d_objectives_deterministic
print 'givens_valid'
print givens_valid
iter_validate = theano.function([], test_objectives, givens=givens_valid)

if config().restart_from_save:
    print 'Load model parameters for resuming'
    resume_metadata = utils.load_pkl(config().restart_from_save)
    nn.layers.set_all_param_values(model.l_out, resume_metadata['param_values'])
    start_chunk_idx = resume_metadata['chunks_since_start'] + 1
    chunk_idxs = range(start_chunk_idx, config().max_nchunks)

    lr = np.float32(utils.current_learning_rate(learning_rate_schedule, start_chunk_idx))
    print '  setting learning rate to %.7f' % lr
    learning_rate.set_value(lr)
    losses_eval_train = resume_metadata['losses_eval_train']
    losses_eval_valid = resume_metadata['losses_eval_valid']
else:
    chunk_idxs = range(config().max_nchunks)
    losses_eval_train = defaultdict(list)
    losses_eval_valid = defaultdict(list)
    start_chunk_idx = 0

train_data_iterator = config().train_data_iterator
valid_data_iterator = config().valid_data_iterator

print
print 'Data'
print 'n train: %d' % train_data_iterator.nsamples
print 'n validation: %d' % valid_data_iterator.nsamples
print 'n chunks per epoch', config().nchunks_per_epoch

print
print 'Train model'
chunk_idx = 0
start_time = time.time()
prev_time = start_time

tmp_losses_train = defaultdict(list)
losses_train_print = defaultdict(list)

# use buffering.buffered_gen_threaded()
for chunk_idx, (x_chunk_train, y_chunk_train, z_chunk_train, id_train) in izip(chunk_idxs, buffering.buffered_gen_threaded(
        train_data_iterator.generate())):
    if chunk_idx in learning_rate_schedule:
        lr = np.float32(learning_rate_schedule[chunk_idx])
        print '  setting learning rate to %.7f' % lr
        print
        learning_rate.set_value(lr)

    # load chunk to GPU
    x_shared.set_value(x_chunk_train)
    y_shared.set_value(y_chunk_train)
    if config().need_enable:
        z_shared.set_value(z_chunk_train)

    # make nbatches_chunk iterations
    for b in xrange(config().nbatches_chunk):
        losses = iter_train(b)
        # print loss
        for obj_idx, obj_name in enumerate(config().order_objectives):
            tmp_losses_train[obj_name].append(losses[obj_idx])
            losses_train_print[obj_name].append(losses[obj_idx])

    if (chunk_idx + 1) % 10 == 0:
        means = []
        for obj_idx, obj_name in enumerate(config().order_objectives):
            mean = np.mean(losses_train_print[obj_name])
            means.append(mean)
            print obj_name, mean
        print 'Chunk %d/%d' % (chunk_idx + 1, config().max_nchunks), sum(means)
        
        losses_train_print = defaultdict(list)

    if ((chunk_idx + 1) % config().validate_every) == 0:
        # calculate mean train loss since the last validation phase
        means = []
        print 'Mean train losses:'
        for obj_idx, obj_name in enumerate(config().order_objectives):
            train_mean = np.mean(tmp_losses_train[obj_name])
            losses_eval_train[obj_name] = train_mean
            means.append(train_mean)
            print obj_name, train_mean
        tmp_losses_train = defaultdict(list)
        print 'Sum of train losses:', sum(means)
        print 'Chunk %d/%d' % (chunk_idx + 1, config().max_nchunks), sum(means)

        # load validation data to GPU
        tmp_losses_valid = defaultdict(list)
        for i, (x_chunk_valid, y_chunk_valid, z_chunk_valid, ids_batch) in enumerate(
                buffering.buffered_gen_threaded(valid_data_iterator.generate(),
                                                buffer_size=2)):
            x_shared.set_value(x_chunk_valid)
            y_shared.set_value(y_chunk_valid)
            if config().need_enable:
                z_shared.set_value(z_chunk_valid)
            losses_valid = iter_validate()
            print i, losses_valid[0], np.sum(losses_valid)
            for obj_idx, obj_name in enumerate(config().order_objectives):
                if z_chunk_valid[0, obj_idx]>0.5:
                    tmp_losses_valid[obj_name].append(losses_valid[obj_idx])


        # calculate validation loss across validation set
        means = [] 
        for obj_idx, obj_name in enumerate(config().order_objectives):
            valid_mean = np.mean(tmp_losses_valid[obj_name])
            losses_eval_valid[obj_name] = valid_mean
            means.append(valid_mean)
            print obj_name, valid_mean
        print 'Sum of mean losses:', sum(means) 


        now = time.time()
        time_since_start = now - start_time
        time_since_prev = now - prev_time
        prev_time = now
        est_time_left = time_since_start * (config().max_nchunks - chunk_idx + 1.) / (chunk_idx + 1. - start_chunk_idx)
        eta = datetime.now() + timedelta(seconds=est_time_left)
        eta_str = eta.strftime("%c")
        print "  %s since start (%.2f s)" % (utils.hms(time_since_start), time_since_prev)
        print "  estimated %s to go (ETA: %s)" % (utils.hms(est_time_left), eta_str)
        print

    if ((chunk_idx + 1) % config().save_every) == 0:
        print
        print 'Chunk %d/%d' % (chunk_idx + 1, config().max_nchunks)
        print 'Saving metadata, parameters'

        with open(metadata_path, 'w') as f:
            pickle.dump({
                'configuration_file': config_name,
                'git_revision_hash': utils.get_git_revision_hash(),
                'experiment_id': expid,
                'chunks_since_start': chunk_idx,
                'losses_eval_train': losses_eval_train,
                'losses_eval_valid': losses_eval_valid,
                'param_values': nn.layers.get_all_param_values(model.l_out)
            }, f, pickle.HIGHEST_PROTOCOL)
            print '  saved to %s' % metadata_path
            print