train.py

#!/usr/bin/env python
# coding: utf-8

import numpy as np
import os
import sys
import glob
import uproot as ur
import matplotlib.pyplot as plt
import time
import seaborn as sns
import tensorflow as tf
from graph_nets import utils_np
from graph_nets import utils_tf
from graph_nets.graphs import GraphsTuple
import sonnet as snt
import argparse
import yaml
import logging
import tensorflow as tf
from tqdm import tqdm

from gn4pions.modules.data import GraphDataGenerator
from gn4pions.modules.models import MultiOutWeightedRegressModel
from gn4pions.modules.utils import convert_to_tuple

sns.set_context('poster')

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("config", default=None, type=str, help="Specify training config file.")
    return parser.parse_args()

if __name__ == "__main__":
    args = get_args()

    # Loading model config
    config_file = args.config
    config = yaml.load(open(config_file), Loader=yaml.FullLoader)

    # Data config
    data_config = config['data']
    data_dir = data_config['data_dir']
    cell_geo_file = data_config['cell_geo_file']
    num_train_files = data_config['num_train_files']
    num_val_files = data_config['num_val_files']
    batch_size = data_config['batch_size']
    shuffle = data_config['shuffle']
    num_procs = data_config['num_procs']
    preprocess = data_config['preprocess']
    output_dir = data_config['output_dir']
    already_preprocessed = data_config['already_preprocessed']

    # Model Config
    model_config = config['model']
    concat_input = model_config['concat_input']

    # Training config
    train_config = config['training']
    epochs = train_config['epochs']
    learning_rate = train_config['learning_rate']
    alpha = train_config['alpha']
    os.environ['CUDA_VISIBLE_DEVICES'] = str(train_config['gpu'])
    physical_devices = tf.config.list_physical_devices('GPU') 
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

    log_freq = train_config['log_freq']
    save_dir = train_config['save_dir'] + config_file.replace('.yaml','').split('/')[-1] + '_' + time.strftime("%Y%m%d")

    os.makedirs(save_dir, exist_ok=True)
    yaml.dump(config, open(save_dir + '/config.yaml', 'w'))

    # Read data and create data generators
    pi0_files = []
    pion_files = np.sort(glob.glob(data_dir+'pion_files/*.npy'))

    train_start = 0
    train_end = train_start + num_train_files
    val_end = train_end + num_val_files

    pi0_train_files = None
    pi0_val_files = None
    pion_train_files = pion_files[train_start:train_end]
    pion_val_files = pion_files[train_end:val_end]

    train_output_dir = None
    val_output_dir = None

    # Get Data
    if preprocess:
        train_output_dir = output_dir + '/train/'
        val_output_dir = output_dir + '/val/'

        if already_preprocessed:
            train_files = np.sort(glob.glob(train_output_dir+'*.p'))[:num_train_files]
            val_files = np.sort(glob.glob(val_output_dir+'*.p'))[:num_val_files]

            pi0_train_files = None
            pi0_val_files = None
            pion_train_files = train_files
            pion_val_files = val_files

            train_output_dir = None
            val_output_dir = None

    # Training Data Generator
    # Will preprocess data if it doesnt find pickled files
    data_gen_train = GraphDataGenerator(pi0_file_list=pi0_train_files,
                                        pion_file_list=pion_train_files,
                                        cellGeo_file=cell_geo_file,
                                        batch_size=batch_size,
                                        shuffle=shuffle,
                                        num_procs=num_procs,
                                        preprocess=preprocess,
                                        output_dir=train_output_dir)

    # Validation Data generator
    # Will preprocess data if it doesnt find pickled files
    data_gen_val = GraphDataGenerator(pi0_file_list=pi0_val_files,
                                      pion_file_list=pion_val_files,
                                      cellGeo_file=cell_geo_file,
                                      batch_size=batch_size,
                                      shuffle=shuffle,
                                      num_procs=num_procs,
                                      preprocess=preprocess,
                                      output_dir=val_output_dir)

    # Get batch of data
    def get_batch(data_iter):
        for graphs, targets in data_iter:
            targets = tf.convert_to_tensor(targets)
            graphs, energies, etas, em_probs, cluster_calib_es, cluster_had_weights, truth_particle_pts, track_pts = convert_to_tuple(graphs)
            yield graphs, targets, energies, etas, em_probs, cluster_calib_es, cluster_had_weights, truth_particle_pts, track_pts

    # Define loss function        
    mae_loss = tf.keras.losses.MeanAbsoluteError()
    bce_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)
    def loss_fn(targets, regress_preds, class_preds):
        regress_loss = mae_loss(targets[:,:1], regress_preds)
        class_loss = bce_loss(targets[:,1:], class_preds)
        combined_loss = alpha*regress_loss + (1 - alpha)*class_loss 
        return regress_loss, class_loss, combined_loss

    # Get a sample graph for tf.function decorator
    samp_graph, samp_target, _, _, _, _, _, _, _ = next(get_batch(data_gen_train.generator()))
    data_gen_train.kill_procs()
    graph_spec = utils_tf.specs_from_graphs_tuple(samp_graph, True, True, True)

    # Training set
    @tf.function(input_signature=[graph_spec, tf.TensorSpec(shape=[None,2], dtype=tf.float32)])
    def train_step(graphs, targets):
        with tf.GradientTape() as tape:
            regress_output, class_output = model(graphs)
            regress_preds = regress_output.globals
            class_preds = class_output.globals
            regress_loss, class_loss, loss = loss_fn(targets, regress_preds, class_preds)

        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
        return regress_loss, class_loss, loss

    # Validation Step
    @tf.function(input_signature=[graph_spec, tf.TensorSpec(shape=[None,2], dtype=tf.float32)])
    def val_step(graphs, targets):
        regress_output, class_output = model(graphs)
        regress_preds = regress_output.globals
        class_preds = class_output.globals
        regress_loss, class_loss, loss = loss_fn(targets, regress_preds, class_preds)
        return regress_loss, class_loss, loss, regress_preds, class_preds

    # Model 
    model = MultiOutWeightedRegressModel(global_output_size=1, num_outputs=2, model_config=model_config)

    # Optimizer
    optimizer = tf.keras.optimizers.Adam(learning_rate)

    # Average epoch losses
    training_loss_epoch = []
    training_loss_regress_epoch = []
    training_loss_class_epoch = []
    val_loss_epoch = []
    val_loss_regress_epoch = []
    val_loss_class_epoch = []

    # Model checkpointing, load latest model if available
    checkpoint = tf.train.Checkpoint(module=model)
    checkpoint_prefix = os.path.join(save_dir, 'latest_model')
    latest = tf.train.latest_checkpoint(save_dir)
    if latest is not None:
        checkpoint.restore(latest)
    else:
        checkpoint.save(checkpoint_prefix)

    # Run training
    curr_loss = 1e5

    for e in range(epochs):

        print(f'\n\nStarting epoch: {e}')
        epoch_start = time.time()

        # Batchwise losses
        training_loss = []
        training_loss_regress = []
        training_loss_class = []
        val_loss = []
        val_loss_regress = []
        val_loss_class = []

        # Train
        print('Training...')
        start = time.time()
        for i, (graph_data_tr, targets_tr, _, _, _, _, _, _, _) in enumerate(get_batch(data_gen_train.generator())):
            losses_tr_rg, losses_tr_cl, losses_tr = train_step(graph_data_tr, targets_tr)

            training_loss.append(losses_tr.numpy())
            training_loss_regress.append(losses_tr_rg.numpy())
            training_loss_class.append(losses_tr_cl.numpy())

            if not (i-1)%log_freq:
                end = time.time()
                print(f'Iter: {i:04d}, ', end='')
                print(f'Tr_loss_mean: {np.mean(training_loss):.4f}, ', end='')
                print(f'Tr_loss_rg_mean: {np.mean(training_loss_regress):.4f}, ', end='') 
                print(f'Tr_loss_cl_mean: {np.mean(training_loss_class):.4f}, ', end='') 
                print(f'Took {end-start:.4f}secs')
                start = time.time()

        training_loss_epoch.append(training_loss)
        training_loss_regress_epoch.append(training_loss_regress)
        training_loss_class_epoch.append(training_loss_class)
        training_end = time.time()

        # validate
        print('\nValidation...')
        all_targets = []
        all_outputs = []
        all_energies = []
        all_etas = []
        all_em_probs = []
        all_cluster_calib_es = []
        all_cluster_had_weights = []
        all_truth_particle_pts = []
        all_track_pts = []
        
        start = time.time()
        for i, (graph_data_val, targets_val, energies_val, etas_val, em_probs_val, cluster_calib_es_val, cluster_had_weights_val, truth_particle_pts_val, track_pts_val) in enumerate(get_batch(data_gen_val.generator())):
            losses_val_rg, losses_val_cl, losses_val, regress_vals, class_vals = val_step(graph_data_val, targets_val)

            targets_val = targets_val.numpy()
            regress_vals = regress_vals.numpy()
            class_vals = class_vals.numpy()

            ### These variables are stored as log_10, so need to exponentiate them again here 
            targets_val[:,0] = 10**targets_val[:,0]
            regress_vals = 10**regress_vals
            class_vals =  tf.math.sigmoid(class_vals)
            energy = 10**graph_data_val.globals 

            output_vals = np.hstack([regress_vals, class_vals])

            val_loss.append(losses_val.numpy())
            val_loss_regress.append(losses_val_rg.numpy())
            val_loss_class.append(losses_val_cl.numpy())

            all_targets.append(targets_val)
            all_outputs.append(output_vals)
            all_energies.append([10**energy for energy in energies_val])
            all_etas.append(etas_val)
            all_em_probs.append(em_probs_val)
            all_cluster_calib_es.append([10**energy for energy in cluster_calib_es_val])
            all_cluster_had_weights.append(cluster_had_weights_val)
            all_truth_particle_pts.append(truth_particle_pts_val)
            all_track_pts.append(track_pts_val)

            if not (i-1)%log_freq:
                end = time.time()
                print(f'Iter: {i:04d}, ', end='')
                print(f'Val_loss_mean: {np.mean(val_loss):.4f}, ', end='')
                print(f'Val_loss_rg_mean: {np.mean(val_loss_regress):.4f}, ', end='') 
                print(f'Val_loss_cl_mean: {np.mean(val_loss_class):.4f}, ', end='') 
                print(f'Took {end-start:.4f}secs')
                start = time.time()

        epoch_end = time.time()

        all_targets = np.concatenate(all_targets)
        all_outputs = np.concatenate(all_outputs)
        all_energies = np.concatenate(all_energies)
        all_etas = np.concatenate(all_etas)
        all_em_probs = np.concatenate(all_em_probs)
        all_cluster_calib_es = np.concatenate(all_cluster_calib_es)
        all_cluster_had_weights = np.concatenate(all_cluster_had_weights)
        all_truth_particle_pts = np.concatenate(all_truth_particle_pts) 
        all_track_pts = np.concatenate(all_track_pts) 
        
        val_loss_epoch.append(val_loss)
        val_loss_regress_epoch.append(val_loss_regress)
        val_loss_class_epoch.append(val_loss_class)

        # Book keeping
        val_mins = int((epoch_end - training_end)/60)
        val_secs = int((epoch_end - training_end)%60)
        training_mins = int((training_end - epoch_start)/60)
        training_secs = int((training_end - epoch_start)%60)
        print(f'\nEpoch {e} ended')
        print(f'Training: {training_mins:2d}:{training_secs:02d}')
        print(f'Validation: {val_mins:2d}:{val_secs:02d}')


        # Save losses
        np.savez(save_dir+'/losses', 
                training=training_loss_epoch, validation=val_loss_epoch,
                training_regress=training_loss_regress_epoch, validation_regress=val_loss_regress_epoch,
                training_class=training_loss_class_epoch, validation_class=val_loss_class_epoch,
                )


        # Checkpoint if validation loss improved
        if np.mean(val_loss)<curr_loss:
            print(f'Loss decreased from {curr_loss:.4f} to {np.mean(val_loss):.4f}')
            print(f'Checkpointing and saving predictions to:\n{save_dir}')
            curr_loss = np.mean(val_loss)
            np.savez(save_dir+'/predictions', 
                    targets=all_targets, 
                    outputs=all_outputs,
                    energies=all_energies,
                    etas=all_etas,
                    em_probs=all_em_probs,
                    cluster_calib_es=all_cluster_calib_es,
                    cluster_had_weights=all_cluster_had_weights,
                    truth_particle_pts=all_truth_particle_pts,
                    track_pts=all_track_pts)
            checkpoint.save(checkpoint_prefix)
        else: 
            print(f'Loss didnt decrease from {curr_loss:.4f}')

        # Decrease learning rate every few epochs
        if not (e+1)%2:   #%20:
            optimizer.learning_rate = optimizer.learning_rate/10
            print(f'Learning rate decreased to: {optimizer.learning_rate.value():.3e}')