training_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import sys\n",
    "import glob\n",
    "import uproot as ur\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "import seaborn as sns\n",
    "import tensorflow as tf\n",
    "from graph_nets import utils_np\n",
    "from graph_nets import utils_tf\n",
    "from graph_nets.graphs import GraphsTuple\n",
    "import sonnet as snt\n",
    "import argparse\n",
    "import yaml\n",
    "import logging\n",
    "import tensorflow as tf\n",
    "from tqdm import tqdm\n",
    "\n",
    "from gn4pions.modules.data import GraphDataGenerator\n",
    "from gn4pions.modules.models import MultiOutWeightedRegressModel\n",
    "from gn4pions.modules.utils import convert_to_tuple\n",
    "\n",
    "sns.set_context('poster')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading model config\n",
    "config_file = 'gn4pions/configs/test.yaml' # for a quick run of the notebook\n",
    "# config_file = 'gn4pions/configs/baseline.yaml' # for actual training\n",
    "config = yaml.load(open(config_file), Loader=yaml.FullLoader)\n",
    "\n",
    "# Data config\n",
    "data_config = config['data']\n",
    "cell_geo_file = data_config['cell_geo_file']\n",
    "data_dir = data_config['data_dir']\n",
    "num_train_files = data_config['num_train_files']\n",
    "num_val_files = data_config['num_val_files']\n",
    "batch_size = data_config['batch_size']\n",
    "shuffle = data_config['shuffle']\n",
    "num_procs = data_config['num_procs']\n",
    "preprocess = data_config['preprocess']\n",
    "output_dir = data_config['output_dir']\n",
    "already_preprocessed = False\n",
    "# already_preprocessed = data_config['already_preprocessed']  # Set to false when running training for first time\n",
    "\n",
    "# Model Config\n",
    "model_config = config['model']\n",
    "concat_input = model_config['concat_input']\n",
    "\n",
    "# Traning Config\n",
    "train_config = config['training']\n",
    "\n",
    "epochs = train_config['epochs']\n",
    "learning_rate = train_config['learning_rate']\n",
    "alpha = train_config['alpha']\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = \"1\"\n",
    "physical_devices = tf.config.list_physical_devices('GPU') \n",
    "tf.config.experimental.set_memory_growth(physical_devices[0], True)\n",
    "log_freq = train_config['log_freq']\n",
    "save_dir = train_config['save_dir'] + config_file.replace('.yaml','').split('/')[-1] + '_' + time.strftime(\"%Y%m%d\")\n",
    "\n",
    "os.makedirs(save_dir, exist_ok=True)\n",
    "yaml.dump(config, open(save_dir + '/config.yaml', 'w'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Preprocessing and saving data to /clusterfs/ml4hep/mpettee/ml4pions/data/onetrack_multicluster/preprocessed//train/\n",
      "Processing file number 0\n",
      "Processing file number 1\n",
      "Processing file number 2\n",
      "Processing file number 3\n",
      "Processing file number 4\n",
      "No pi0 files.\n",
      "No pi0 files.\n",
      "No pi0 files.\n",
      "No pi0 files.\n",
      "No pi0 files.\n",
      "Finished processing 4 files\n",
      "Finished processing 2 files\n",
      "Finished processing 3 files\n",
      "Finished processing 1 files\n",
      "Finished processing 0 files\n",
      "\n",
      "Preprocessing and saving data to /clusterfs/ml4hep/mpettee/ml4pions/data/onetrack_multicluster/preprocessed//val/\n",
      "Processing file number 0\n",
      "Processing file number 1\n",
      "Processing file number 2\n",
      "Processing file number 3\n",
      "Processing file number 4\n",
      "No pi0 files.\n",
      "No pi0 files.\n",
      "No pi0 files.\n",
      "No pi0 files.\n",
      "No pi0 files.\n",
      "Finished processing 1 files\n",
      "Finished processing 0 files\n",
      "Finished processing 3 files\n",
      "Finished processing 4 files\n",
      "Finished processing 2 files\n"
     ]
    }
   ],
   "source": [
    "# Read data and create data generators\n",
    "\n",
    "pi0_files = np.sort(glob.glob(data_dir+'pi0_files/*.npy'))\n",
    "pion_files = np.sort(glob.glob(data_dir+'pion_files/*.npy'))\n",
    "\n",
    "len(pi0_files)\n",
    "len(pion_files)\n",
    "\n",
    "train_start = 0\n",
    "train_end = train_start + num_train_files\n",
    "val_end = train_end + num_val_files\n",
    "\n",
    "pi0_train_files = pi0_files[train_start:train_end]\n",
    "pi0_val_files = pi0_files[train_end:val_end]\n",
    "pion_train_files = pion_files[train_start:train_end]\n",
    "pion_val_files = pion_files[train_end:val_end]\n",
    "\n",
    "train_output_dir = None\n",
    "val_output_dir = None\n",
    "\n",
    "# Get Data\n",
    "if preprocess:\n",
    "    train_output_dir = output_dir + '/train/'\n",
    "    val_output_dir = output_dir + '/val/'\n",
    "\n",
    "    if already_preprocessed:\n",
    "        train_files = np.sort(glob.glob(train_output_dir+'*.p'))[:num_train_files]\n",
    "        val_files = np.sort(glob.glob(val_output_dir+'*.p'))[:num_val_files]\n",
    "\n",
    "        pi0_train_files = None\n",
    "        pi0_val_files = None\n",
    "        pion_train_files = train_files\n",
    "        pion_val_files = val_files\n",
    "\n",
    "        train_output_dir = None\n",
    "        val_output_dir = None\n",
    "\n",
    "# Traning Data Generator\n",
    "# Will preprocess data if it doesnt find pickled files\n",
    "data_gen_train = GraphDataGenerator(pi0_file_list=pi0_train_files,\n",
    "                                    pion_file_list=pion_train_files,\n",
    "                                    cellGeo_file=cell_geo_file,\n",
    "                                    batch_size=batch_size,\n",
    "                                    shuffle=shuffle,\n",
    "                                    num_procs=num_procs,\n",
    "                                    preprocess=preprocess,\n",
    "                                    output_dir=train_output_dir)\n",
    "\n",
    "# Validation Data generator\n",
    "# Will preprocess data if it doesnt find pickled files\n",
    "data_gen_val = GraphDataGenerator(pi0_file_list=pi0_val_files,\n",
    "                                  pion_file_list=pion_val_files,\n",
    "                                  cellGeo_file=cell_geo_file,\n",
    "                                  batch_size=batch_size,\n",
    "                                  shuffle=shuffle,\n",
    "                                  num_procs=num_procs,\n",
    "                                  preprocess=preprocess,\n",
    "                                  output_dir=val_output_dir)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get batch of data\n",
    "def get_batch(data_iter):\n",
    "    for graphs, targets in data_iter:\n",
    "        graphs, energies, etas = convert_to_tuple(graphs)\n",
    "        targets = tf.convert_to_tensor(targets)\n",
    "        yield graphs, targets, energies, etas\n",
    "        \n",
    "# Define loss function        \n",
    "mae_loss = tf.keras.losses.MeanAbsoluteError()\n",
    "bce_loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)\n",
    "def loss_fn(targets, regress_preds, class_preds):\n",
    "    regress_loss = mae_loss(targets[:,:1], regress_preds)\n",
    "    class_loss = bce_loss(targets[:,1:], class_preds)\n",
    "    combined_loss = alpha*regress_loss + (1 - alpha)*class_loss \n",
    "    return regress_loss, class_loss, combined_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get a sample graph for tf.function decorator\n",
    "samp_graph, samp_target, samp_e, samp_eta = next(get_batch(data_gen_train.generator()))\n",
    "data_gen_train.kill_procs()\n",
    "graph_spec = utils_tf.specs_from_graphs_tuple(samp_graph, True, True, True)\n",
    "\n",
    "# Training set\n",
    "@tf.function(input_signature=[graph_spec, tf.TensorSpec(shape=[None,2], dtype=tf.float32)])\n",
    "def train_step(graphs, targets):\n",
    "    with tf.GradientTape() as tape:\n",
    "        regress_output, class_output = model(graphs)\n",
    "        regress_preds = regress_output.globals\n",
    "        class_preds = class_output.globals\n",
    "        regress_loss, class_loss, loss = loss_fn(targets, regress_preds, class_preds)\n",
    "\n",
    "    gradients = tape.gradient(loss, model.trainable_variables)\n",
    "    optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "    return regress_loss, class_loss, loss\n",
    "\n",
    "# Validation Step\n",
    "@tf.function(input_signature=[graph_spec, tf.TensorSpec(shape=[None,2], dtype=tf.float32)])\n",
    "def val_step(graphs, targets):\n",
    "    regress_output, class_output = model(graphs)\n",
    "    regress_preds = regress_output.globals\n",
    "    class_preds = class_output.globals\n",
    "    regress_loss, class_loss, loss = loss_fn(targets, regress_preds, class_preds)\n",
    "    return regress_loss, class_loss, loss, regress_preds, class_preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model \n",
    "model = MultiOutWeightedRegressModel(global_output_size=1, num_outputs=2, model_config=model_config)\n",
    "\n",
    "# Optimizer\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate)\n",
    "\n",
    "# Average epoch losses\n",
    "training_loss_epoch = []\n",
    "training_loss_regress_epoch = []\n",
    "training_loss_class_epoch = []\n",
    "val_loss_epoch = []\n",
    "val_loss_regress_epoch = []\n",
    "val_loss_class_epoch = []\n",
    "\n",
    "# Model checkpointing, load latest model if available\n",
    "checkpoint = tf.train.Checkpoint(module=model)\n",
    "checkpoint_prefix = os.path.join(save_dir, 'latest_model')\n",
    "latest = tf.train.latest_checkpoint(save_dir)\n",
    "if latest is not None:\n",
    "    checkpoint.restore(latest)\n",
    "else:\n",
    "    checkpoint.save(checkpoint_prefix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for i, (graph_data_tr, targets_tr, energies_tr, etas_tr) in enumerate(get_batch(data_gen_train.generator())):\n",
    "#     print(\"Targets: \",targets_tr)\n",
    "#     print(\"graph\", graph_data_tr)\n",
    "#     print(\"Energies\", energies_tr)\n",
    "#     print(\"etas\", etas_tr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run training\n",
    "curr_loss = 1e5\n",
    "\n",
    "epochs = 1\n",
    "\n",
    "for e in range(epochs):\n",
    "\n",
    "    print(f'\\n\\nStarting epoch: {e}')\n",
    "    epoch_start = time.time()\n",
    "    \n",
    "    # Batchwise losses\n",
    "    training_loss = []\n",
    "    training_loss_regress = []\n",
    "    training_loss_class = []\n",
    "    val_loss = []\n",
    "    val_loss_regress = []\n",
    "    val_loss_class = []\n",
    "\n",
    "    # Train\n",
    "    print('Training...')\n",
    "    start = time.time()\n",
    "    for i, (graph_data_tr, targets_tr, energies_tr, etas_tr) in enumerate(get_batch(data_gen_train.generator())):\n",
    "        losses_tr_rg, losses_tr_cl, losses_tr = train_step(graph_data_tr, targets_tr)\n",
    "\n",
    "        training_loss.append(losses_tr.numpy())\n",
    "        training_loss_regress.append(losses_tr_rg.numpy())\n",
    "        training_loss_class.append(losses_tr_cl.numpy())\n",
    "\n",
    "        if not (i-1)%log_freq:\n",
    "            end = time.time()\n",
    "            print(f'Iter: {i:04d}, ', end='')\n",
    "            print(f'Tr_loss_mean: {np.mean(training_loss):.4f}, ', end='')\n",
    "            print(f'Tr_loss_rg_mean: {np.mean(training_loss_regress):.4f}, ', end='') \n",
    "            print(f'Tr_loss_cl_mean: {np.mean(training_loss_class):.4f}, ', end='') \n",
    "            print(f'Took {end-start:.4f}secs')\n",
    "            start = time.time()\n",
    "                  \n",
    "    training_loss_epoch.append(training_loss)\n",
    "    training_loss_regress_epoch.append(training_loss_regress)\n",
    "    training_loss_class_epoch.append(training_loss_class)\n",
    "    training_end = time.time()\n",
    "\n",
    "    # validate\n",
    "    print('\\nValidation...')\n",
    "    all_targets = []\n",
    "    all_outputs = []\n",
    "    all_energies = []\n",
    "    all_etas = []\n",
    "    start = time.time()\n",
    "    for i, (graph_data_val, targets_val, energies_val, etas_val) in enumerate(get_batch(data_gen_val.generator())):\n",
    "        losses_val_rg, losses_val_cl, losses_val, regress_vals, class_vals = val_step(graph_data_val, targets_val)\n",
    "\n",
    "        targets_val = targets_val.numpy()\n",
    "        regress_vals = regress_vals.numpy()\n",
    "        class_vals = class_vals.numpy()\n",
    "\n",
    "        targets_val[:,0] = 10**targets_val[:,0]\n",
    "        regress_vals = 10**regress_vals\n",
    "        class_vals =  tf.math.sigmoid(class_vals)\n",
    "\n",
    "        output_vals = np.hstack([regress_vals, class_vals])\n",
    "        \n",
    "        val_loss.append(losses_val.numpy())\n",
    "        val_loss_regress.append(losses_val_rg.numpy())\n",
    "        val_loss_class.append(losses_val_cl.numpy())\n",
    "\n",
    "        all_targets.append(targets_val)\n",
    "        all_outputs.append(output_vals)\n",
    "        all_energies.append([10**energy for energy in energies_val])\n",
    "        all_etas.append(etas_val)\n",
    "\n",
    "        if not (i-1)%log_freq:\n",
    "            end = time.time()\n",
    "            print(f'Iter: {i:04d}, ', end='')\n",
    "            print(f'Val_loss_mean: {np.mean(val_loss):.4f}, ', end='')\n",
    "            print(f'Val_loss_rg_mean: {np.mean(val_loss_regress):.4f}, ', end='') \n",
    "            print(f'Val_loss_cl_mean: {np.mean(val_loss_class):.4f}, ', end='') \n",
    "            print(f'Took {end-start:.4f}secs')\n",
    "            start = time.time()\n",
    "\n",
    "    epoch_end = time.time()\n",
    "\n",
    "    all_targets = np.concatenate(all_targets)\n",
    "    all_outputs = np.concatenate(all_outputs)    \n",
    "    all_energies = np.concatenate(all_energies)\n",
    "    all_etas = np.concatenate(all_etas)\n",
    "    \n",
    "    val_loss_epoch.append(val_loss)\n",
    "    val_loss_regress_epoch.append(val_loss_regress)\n",
    "    val_loss_class_epoch.append(val_loss_class)\n",
    "\n",
    "    \n",
    "    # Book keeping\n",
    "    val_mins = int((epoch_end - training_end)/60)\n",
    "    val_secs = int((epoch_end - training_end)%60)\n",
    "    training_mins = int((training_end - epoch_start)/60)\n",
    "    training_secs = int((training_end - epoch_start)%60)\n",
    "    print(f'\\nEpoch {e} ended')\n",
    "    print(f'Training: {training_mins:2d}:{training_secs:02d}')\n",
    "    print(f'Validation: {val_mins:2d}:{val_secs:02d}')\n",
    "    \n",
    "    \n",
    "    # Save losses\n",
    "    np.savez(save_dir+'/losses', \n",
    "            training=training_loss_epoch, validation=val_loss_epoch,\n",
    "            training_regress=training_loss_regress_epoch, validation_regress=val_loss_regress_epoch,\n",
    "            training_class=training_loss_class_epoch, validation_class=val_loss_class_epoch,\n",
    "            )\n",
    "\n",
    "    \n",
    "    # Checkpoint if validation loss improved\n",
    "    if np.mean(val_loss)<curr_loss:\n",
    "        print(f'Loss decreased from {curr_loss:.4f} to {np.mean(val_loss):.4f}')\n",
    "        print(f'Checkpointing and saving predictions to:\\n{save_dir}')\n",
    "        curr_loss = np.mean(val_loss)\n",
    "        np.savez(save_dir+'/predictions', \n",
    "                targets=all_targets, \n",
    "                outputs=all_outputs,\n",
    "                energies=all_energies,\n",
    "                etas=all_etas)\n",
    "        checkpoint.save(checkpoint_prefix)\n",
    "    else: \n",
    "        print(f'Loss didnt decrease from {curr_loss:.4f}')\n",
    "    \n",
    "    \n",
    "    # Decrease learning rate every few epochs\n",
    "    if not (e+1)%2:   #%20:\n",
    "        optimizer.learning_rate = optimizer.learning_rate/10\n",
    "        print(f'Learning rate decreased to: {optimizer.learning_rate.value():.3e}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nbdev",
   "language": "python",
   "name": "nbdev"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}