tutorials/T03-ModelLoading.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fc706538-3cc8-4a91-ad12-da97f708f0b5",
   "metadata": {},
   "source": [
    "# Tutorial 03 - Model Loading\n",
    "\n",
    "This tutorial demonstrates how to create tensorflow models for supervised and weakly trainings using the paws API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7b3d88c4-5706-4b45-9403-c908343a68d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from paws.components import ModelLoader\n",
    "\n",
    "# here outdir is the directory to which model outputs are saved, this is only needed for loading\n",
    "# semi-weakly models which is built from supervised models (it needs to know where to load the\n",
    "# supervised models)\n",
    "outdir = \"/pscratch/sd/c/chlcheng/projects/paws/outputs\"\n",
    "\n",
    "# for more details of the parameters, use help(ModelLoader.__init__)\n",
    "model_loader = ModelLoader(feature_level=\"high_level\", decay_modes='qq,qqq',\n",
    "                           variables=\"3,5,6\", noise_dimension=0, outdir=outdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8dea3454-565f-4019-928c-95f90c0e43c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-06-01 21:33:56.967190: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2024-06-01 21:33:56.968193: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2024-06-01 21:33:57.070216: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2024-06-01 21:33:57.307532: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-06-01 21:34:00.007541: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
      "2024-06-01 21:34:04.050155: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1929] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 22565 MB memory:  -> device: 0, name: NVIDIA A100-PCIE-40GB, pci bus id: 0000:c3:00.0, compute capability: 8.0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] Number of events in each dataset split:\n",
      "            train     val    test\n",
      "QCD         60700   30327   30325\n",
      "extra_QCD  306450  153208  153200\n",
      "W_qq        37250   18622   18600\n",
      "W_qqq       38111   19050   19050\n",
      "total      442511  221207  221175\n",
      "num_batch     433     217     216\n"
     ]
    }
   ],
   "source": [
    "from paws.components import DataLoader\n",
    "\n",
    "datadir = \"/pscratch/sd/c/chlcheng/projects/paws/datasets\"\n",
    "\n",
    "# data loading is needed to get the shape and dtype of the train features\n",
    "data_loader = DataLoader(datadir=datadir, feature_level=\"high_level\", decay_modes='qq,qqq',\n",
    "                         variables=\"3,5,6\", noise_dimension=0, seed=2023)\n",
    "datasets = data_loader.get_datasets(mass_point=[300, 300])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a7050448-1a54-4fdc-97ca-a92c24bafca0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'jet_features': {'shape': [2, 7], 'dtype': 'float64'},\n",
       " 'param_masses': {'shape': [2], 'dtype': 'float64'},\n",
       " 'label': {'shape': [1], 'dtype': 'int64'}}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "feature_metadata = data_loader.feature_metadata\n",
    "feature_metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "512e3319-bfce-43be-93e5-f963093bc70b",
   "metadata": {},
   "source": [
    "### Dedicated supervised model (also used for ideal weakly training)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bf0dc135-236e-4e9c-8e91-442539601fa3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"HighLevel\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                Output Shape                 Param #   Connected to                  \n",
      "==================================================================================================\n",
      " jet_features (InputLayer)   [(None, 2, 3)]               0         []                            \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape (TFOpLa  (3,)                         0         ['jet_features[0][0]']        \n",
      " mbda)                                                                                            \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem (  (2,)                         0         ['tf.compat.v1.shape[0][0]']  \n",
      " SlicingOpLambda)                                                                                 \n",
      "                                                                                                  \n",
      " tf.math.reduce_prod (TFOpL  ()                           0         ['tf.__operators__.getitem[0][\n",
      " ambda)                                                             0]']                          \n",
      "                                                                                                  \n",
      " tf.reshape (TFOpLambda)     (None, 6)                    0         ['jet_features[0][0]',        \n",
      "                                                                     'tf.math.reduce_prod[0][0]'] \n",
      "                                                                                                  \n",
      " dense (Dense)               (None, 256)                  1792      ['tf.reshape[0][0]']          \n",
      "                                                                                                  \n",
      " dense_1 (Dense)             (None, 128)                  32896     ['dense[0][0]']               \n",
      "                                                                                                  \n",
      " dense_2 (Dense)             (None, 64)                   8256      ['dense_1[0][0]']             \n",
      "                                                                                                  \n",
      " dense_3 (Dense)             (None, 1)                    65        ['dense_2[0][0]']             \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 43009 (168.00 KB)\n",
      "Trainable params: 43009 (168.00 KB)\n",
      "Non-trainable params: 0 (0.00 Byte)\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "fs_model_dedicated = model_loader.get_supervised_model(feature_metadata, parametric=False)\n",
    "fs_model_dedicated.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e224019-4cfe-4262-ae75-2d1b0786f7c9",
   "metadata": {},
   "source": [
    "### Parameterised supervised model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "55960df3-2df2-4ea3-86ce-4d8c6dd15b83",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"HighLevel\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                Output Shape                 Param #   Connected to                  \n",
      "==================================================================================================\n",
      " param_masses (InputLayer)   [(None, 2)]                  0         []                            \n",
      "                                                                                                  \n",
      " jet_features (InputLayer)   [(None, 2, 3)]               0         []                            \n",
      "                                                                                                  \n",
      " tf.expand_dims (TFOpLambda  (None, 2, 1)                 0         ['param_masses[0][0]']        \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.concat (TFOpLambda)      (None, 2, 4)                 0         ['jet_features[0][0]',        \n",
      "                                                                     'tf.expand_dims[0][0]']      \n",
      "                                                                                                  \n",
      " tf.compat.v1.shape_1 (TFOp  (3,)                         0         ['tf.concat[0][0]']           \n",
      " Lambda)                                                                                          \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_1  (2,)                         0         ['tf.compat.v1.shape_1[0][0]']\n",
      "  (SlicingOpLambda)                                                                               \n",
      "                                                                                                  \n",
      " tf.math.reduce_prod_1 (TFO  ()                           0         ['tf.__operators__.getitem_1[0\n",
      " pLambda)                                                           ][0]']                        \n",
      "                                                                                                  \n",
      " tf.reshape_1 (TFOpLambda)   (None, 8)                    0         ['tf.concat[0][0]',           \n",
      "                                                                     'tf.math.reduce_prod_1[0][0]'\n",
      "                                                                    ]                             \n",
      "                                                                                                  \n",
      " dense_4 (Dense)             (None, 256)                  2304      ['tf.reshape_1[0][0]']        \n",
      "                                                                                                  \n",
      " dense_5 (Dense)             (None, 128)                  32896     ['dense_4[0][0]']             \n",
      "                                                                                                  \n",
      " dense_6 (Dense)             (None, 64)                   8256      ['dense_5[0][0]']             \n",
      "                                                                                                  \n",
      " dense_7 (Dense)             (None, 1)                    65        ['dense_6[0][0]']             \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 43521 (170.00 KB)\n",
      "Trainable params: 43521 (170.00 KB)\n",
      "Non-trainable params: 0 (0.00 Byte)\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "fs_model_parameterised = model_loader.get_supervised_model(feature_metadata, parametric=True)\n",
    "fs_model_parameterised.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e324a1a-c17d-497f-9ae2-049d2abeed60",
   "metadata": {},
   "source": [
    "### Semi-weakly model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5e69792a-4f29-4e3f-ba4d-e82eb7536042",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"SemiWeakly\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                Output Shape                 Param #   Connected to                  \n",
      "==================================================================================================\n",
      " jet_features (InputLayer)   [(None, 2, 3)]               0         []                            \n",
      "                                                                                                  \n",
      " tf.ones_like (TFOpLambda)   (None, 2, 3)                 0         ['jet_features[0][0]']        \n",
      "                                                                                                  \n",
      " tf.ones_like_1 (TFOpLambda  (None, 2, 3)                 0         ['jet_features[0][0]']        \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_2  (None,)                      0         ['tf.ones_like[0][0]']        \n",
      "  (SlicingOpLambda)                                                                               \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_3  (None,)                      0         ['tf.ones_like_1[0][0]']      \n",
      "  (SlicingOpLambda)                                                                               \n",
      "                                                                                                  \n",
      " model (Functional)          (None, 1)                    1         ['tf.__operators__.getitem_2[0\n",
      "                                                                    ][0]']                        \n",
      "                                                                                                  \n",
      " model_1 (Functional)        (None, 1)                    1         ['tf.__operators__.getitem_3[0\n",
      "                                                                    ][0]']                        \n",
      "                                                                                                  \n",
      " tf.math.multiply (TFOpLamb  (None, 1)                    0         ['model[0][0]']               \n",
      " da)                                                                                              \n",
      "                                                                                                  \n",
      " tf.math.multiply_1 (TFOpLa  (None, 1)                    0         ['model_1[0][0]']             \n",
      " mbda)                                                                                            \n",
      "                                                                                                  \n",
      " concatenate (Concatenate)   (None, 2)                    0         ['tf.math.multiply[0][0]',    \n",
      "                                                                     'tf.math.multiply_1[0][0]']  \n",
      "                                                                                                  \n",
      " tf.ones_like_3 (TFOpLambda  (None, 2, 3)                 0         ['jet_features[0][0]']        \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " HighLevel_3prong (Function  (None, 1)                    43521     ['jet_features[0][0]',        \n",
      " al)                                                                 'concatenate[0][0]']         \n",
      "                                                                                                  \n",
      " HighLevel_2prong (Function  (None, 1)                    43521     ['jet_features[0][0]',        \n",
      " al)                                                                 'concatenate[0][0]']         \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_5  (None,)                      0         ['tf.ones_like_3[0][0]']      \n",
      "  (SlicingOpLambda)                                                                               \n",
      "                                                                                                  \n",
      " tf.math.subtract_1 (TFOpLa  (None, 1)                    0         ['HighLevel_3prong[0][0]']    \n",
      " mbda)                                                                                            \n",
      "                                                                                                  \n",
      " tf.math.subtract (TFOpLamb  (None, 1)                    0         ['HighLevel_2prong[0][0]']    \n",
      " da)                                                                                              \n",
      "                                                                                                  \n",
      " model_3 (Functional)        (None, 1)                    1         ['tf.__operators__.getitem_5[0\n",
      "                                                                    ][0]']                        \n",
      "                                                                                                  \n",
      " tf.__operators__.add_1 (TF  (None, 1)                    0         ['tf.math.subtract_1[0][0]']  \n",
      " OpLambda)                                                                                        \n",
      "                                                                                                  \n",
      " tf.__operators__.add (TFOp  (None, 1)                    0         ['tf.math.subtract[0][0]']    \n",
      " Lambda)                                                                                          \n",
      "                                                                                                  \n",
      " tf.math.truediv_1 (TFOpLam  (None, 1)                    0         ['HighLevel_3prong[0][0]',    \n",
      " bda)                                                                'tf.__operators__.add_1[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      " tf.math.subtract_2 (TFOpLa  (None, 1)                    0         ['model_3[0][0]']             \n",
      " mbda)                                                                                            \n",
      "                                                                                                  \n",
      " tf.math.truediv (TFOpLambd  (None, 1)                    0         ['HighLevel_2prong[0][0]',    \n",
      " a)                                                                  'tf.__operators__.add[0][0]']\n",
      "                                                                                                  \n",
      " tf.ones_like_2 (TFOpLambda  (None, 2, 3)                 0         ['jet_features[0][0]']        \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.multiply_2 (TFOpLa  (None, 1)                    0         ['model_3[0][0]',             \n",
      " mbda)                                                               'tf.math.truediv_1[0][0]']   \n",
      "                                                                                                  \n",
      " tf.math.multiply_3 (TFOpLa  (None, 1)                    0         ['tf.math.subtract_2[0][0]',  \n",
      " mbda)                                                               'tf.math.truediv[0][0]']     \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem_4  (None,)                      0         ['tf.ones_like_2[0][0]']      \n",
      "  (SlicingOpLambda)                                                                               \n",
      "                                                                                                  \n",
      " tf.__operators__.add_2 (TF  (None, 1)                    0         ['tf.math.multiply_2[0][0]',  \n",
      " OpLambda)                                                           'tf.math.multiply_3[0][0]']  \n",
      "                                                                                                  \n",
      " model_2 (Functional)        (None, 1)                    1         ['tf.__operators__.getitem_4[0\n",
      "                                                                    ][0]']                        \n",
      "                                                                                                  \n",
      " tf.math.subtract_3 (TFOpLa  (None, 1)                    0         ['tf.__operators__.add_2[0][0]\n",
      " mbda)                                                              ']                            \n",
      "                                                                                                  \n",
      " tf.math.multiply_4 (TFOpLa  (None, 1)                    0         ['model_2[0][0]',             \n",
      " mbda)                                                               'tf.math.subtract_3[0][0]']  \n",
      "                                                                                                  \n",
      " tf.__operators__.add_3 (TF  (None, 1)                    0         ['tf.math.multiply_4[0][0]']  \n",
      " OpLambda)                                                                                        \n",
      "                                                                                                  \n",
      " tf.__operators__.add_4 (TF  (None, 1)                    0         ['tf.__operators__.add_3[0][0]\n",
      " OpLambda)                                                          ']                            \n",
      "                                                                                                  \n",
      " tf.math.truediv_2 (TFOpLam  (None, 1)                    0         ['tf.__operators__.add_3[0][0]\n",
      " bda)                                                               ',                            \n",
      "                                                                     'tf.__operators__.add_4[0][0]\n",
      "                                                                    ']                            \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 87046 (340.02 KB)\n",
      "Trainable params: 4 (16.00 Byte)\n",
      "Non-trainable params: 87042 (340.01 KB)\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# get paths to the supervised models\n",
    "fs_ckpt_dir = model_loader.get_checkpoint_dir(\"param_supervised\", decay_mode=\"qq\", split_index=0, version=\"v1\")\n",
    "fs_model_path = os.path.join(fs_ckpt_dir, model_loader.path_manager.get_basename('model_full_train'))\n",
    "fs_ckpt_dir_2 = model_loader.get_checkpoint_dir(\"param_supervised\", decay_mode=\"qqq\", split_index=0, version=\"v1\")\n",
    "fs_model_path_2 = os.path.join(fs_ckpt_dir_2, model_loader.path_manager.get_basename('model_full_train'))\n",
    "\n",
    "ws_model = model_loader.get_semi_weakly_model(feature_metadata, fs_model_path=fs_model_path,\n",
    "                                              fs_model_path_2=fs_model_path_2)\n",
    "ws_model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd8c6b00-315a-41f8-8209-0e88ffba180a",
   "metadata": {},
   "source": [
    "### Train setup and configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c5c61a2b-2615-4ee6-a91a-d62863476bcf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'loss': <aliad.interface.tensorflow.losses.ScaledBinaryCrossentropy at 0x7fb72878f390>,\n",
       " 'metrics': ['accuracy'],\n",
       " 'epochs': 100,\n",
       " 'optimizer': 'Adam',\n",
       " 'optimizer_config': {'learning_rate': 0.05,\n",
       "  'clipvalue': 0.0001,\n",
       "  'clipnorm': 0.0001},\n",
       " 'checkpoint_dir': '/pscratch/sd/c/chlcheng/projects/paws/outputs/train_results/semi_weakly/high_level/qq_qqq/300_300/SR_var_3_5_6_noise_0_v1/mu_0p005_alpha_0p5/split_0/trial_0',\n",
       " 'callbacks': {'lr_scheduler': {'initial_lr': 0.05,\n",
       "   'lr_decay_factor': 0.5,\n",
       "   'patience': 5,\n",
       "   'min_lr': 1e-06,\n",
       "   'verbose': True},\n",
       "  'early_stopping': {'monitor': 'val_loss',\n",
       "   'patience': 20,\n",
       "   'restore_best_weights': True},\n",
       "  'model_checkpoint': {'save_weights_only': True, 'save_freq': 'epoch'},\n",
       "  'metrics_logger': {'save_freq': -1},\n",
       "  'weights_logger': {'save_freq': -1, 'display_weight': True}}}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get path to checkpoint directory based on the desired scenario\n",
    "ws_ckpt_dir = model_loader.get_checkpoint_dir(\"semi_weakly\", split_index=0, version=\"v1\",\n",
    "                                              mass_point=[300, 300], mu=0.005, alpha=0.5,\n",
    "                                              trial=0)\n",
    "# get train configuration\n",
    "train_config = model_loader.get_train_config(checkpoint_dir=ws_ckpt_dir, model_type='semi_weakly')\n",
    "train_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "47eeb8df-a905-488f-8748-4fa2502a4616",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'early_stopping': <aliad.interface.tensorflow.callbacks.EarlyStopping at 0x7fb7287c9590>,\n",
       " 'model_checkpoint': <keras.src.callbacks.ModelCheckpoint at 0x7fb7287caad0>,\n",
       " 'metrics_logger': <aliad.interface.tensorflow.callbacks.MetricsLogger at 0x7fb7287cae90>,\n",
       " 'lr_scheduler': <aliad.interface.tensorflow.callbacks.LearningRateScheduler at 0x7fb7287c9f90>,\n",
       " 'weights_logger': <aliad.interface.tensorflow.callbacks.WeightsLogger at 0x7fb7287ca310>}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get list of callbacks used in training\n",
    "callbacks = model_loader.get_callbacks(model_type='semi_weakly', config=train_config)\n",
    "callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c17470c8-53ad-4444-9ca1-a1d032c8859c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compile model\n",
    "model_loader.compile_model(ws_model, train_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "28d15ca9-5ec3-4670-8abf-21a458ea7271",
   "metadata": {},
   "outputs": [],
   "source": [
    "# manually set trainable parameters in the semi-weakly model\n",
    "# note the unit of m1 and m2 here is 10GeV, mu is the natural log of actual signal fraction\n",
    "import numpy as np\n",
    "model_loader.set_semi_weakly_model_weights(ws_model, m1=0.5, m2=0.5, mu=np.log(1e-3), alpha=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "969cf8e9-621d-4163-a402-b314ae05793f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'m1': 0.5, 'm2': 0.5, 'alpha': 0.5, 'mu': -6.9077554}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_loader.get_semi_weakly_model_weights(ws_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2891d3d-7feb-4133-b48d-6d11a69f2fd9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}