stt3_gnn_infer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## _Inference_\n",
    "\n",
    "**_Inference_** is done using callbacks defined in the _LightningModules/GNN/Models/inference.py_. The callbacks run during the _test_step()_ _a.k.a_ model _**evalution**_.\n",
    "\n",
    "\n",
    "### _How to Run Inference?_\n",
    "\n",
    "1. _`traintrack config/pipeline_quickstart.yaml`_: One can use `--inference` flag to run only the `test_step()` (don't forget to give `resume_id` of a checkpoint)\n",
    "2. _`infer.ipynb`_ notebook runs the _pl.Trainer().test()_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os, glob, yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pprint\n",
    "from tqdm import tqdm\n",
    "import trackml.dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchmetrics\n",
    "import pytorch_lightning as pl\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.loader import DataLoader\n",
    "import itertools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['EXATRKX_DATA'] = os.path.abspath(os.curdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from LightningModules.GNN import InteractionGNN\n",
    "from LightningModules.GNN import GNNBuilder, GNNMetrics\n",
    "from LightningModules.GNN.Models.infer import GNNTelemetry"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HfvKUCvjUlQA"
   },
   "source": [
    "### _Load Checkpoint_\n",
    "\n",
    "Lightning automatically saves a checkpoint for you in your current working directory, with the state of your last training epoch. We have checkpoint stored after training is finished.\n",
    "\n",
    "```python\n",
    "# load a LightningModule along with its weights & hyperparameters from a checkpoint\n",
    "model = MyLightningModule.load_from_checkpoint(\"/path/to/checkpoint.ckpt\")\n",
    "print(model.input_dir)\n",
    "```\n",
    "\n",
    "Note that we have saved our hyperparameters when our **LightningModule** was initialized i.e. `self.save_hyperparameters(hparams)`\n",
    "\n",
    "```python\n",
    "# hyperparameters are saved to the “hyper_parameters” key in the checkpoint, to access them\n",
    "checkpoint = torch.load(path/to/checkpoint, map_location=device)\n",
    "print(checkpoint[\"hyper_parameters\"])\n",
    "```\n",
    "\n",
    "One can also initialize the model with different hyperparameters (if they are saved).\n",
    "\n",
    "\n",
    "For more details, consult [Lighting Checkpointing](https://pytorch-lightning.readthedocs.io/en/stable/common/checkpointing.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _Get Checkpoint Hparams_\n",
    "\n",
    "- Either from the configs folder \n",
    "- Or extract it from the checkpoint, favoured if model is trained and evaluated on two different machines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load processing config file (trusted source)\n",
    "config = None\n",
    "config_file = os.path.join(os.curdir, 'LightningModules/GNN/configs/train_alldata_IGNN.yaml')\n",
    "with open(config_file) as f:\n",
    "    try:\n",
    "        config = yaml.load(f, Loader=yaml.FullLoader) # equiv: yaml.full_load(f)\n",
    "    except yaml.YAMLError as e:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Model Checkpoint\n",
    "# ckpnt_path = \"run_all/lightning_models/lightning_checkpoints/GNNStudy/a58b2mlx/checkpoints/last.ckpt\"  # muons with Filtering=True\n",
    "# ckpnt_path = \"run_all/lightning_models/lightning_checkpoints/HypGNN/uibb0ir9/checkpoints/last.ckpt\"  # fwp with Filtering=True\n",
    "ckpnt_path = \"run_all/lightning_models/lightning_checkpoints/HypGNN/p9rjmknq/checkpoints/last.ckpt\"  # fwp with Filtering=False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = torch.load(ckpnt_path, map_location=device)\n",
    "hparams = checkpoint[\"hyper_parameters\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View Hyperparameters\n",
    "print(hparams)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# One Can Modify Hyperparameters\n",
    "hparams[\"checkpoint_path\"] = ckpnt_path\n",
    "hparams[\"input_dir\"] = \"run_quick/fwp_feature_store\"\n",
    "hparams[\"output_dir\"] = \"run_quick/fwp_gnn_processed\"\n",
    "hparams[\"artifact_library\"] = \"lightning_models/lightning_checkpoints\"\n",
    "hparams[\"train_split\"] = [0, 0, 19803]\n",
    "hparams[\"map_location\"] = device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View Hyperparameters (Modified)\n",
    "print(hparams)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _Get Checkpoint Model_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Init EdgeClassifier with New Config\n",
    "model = InteractionGNN(hparams)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.hparams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Checkpoint with New Config (It will Provide Path and Other Parameters, Most will be Overwritten)\n",
    "model = model.load_from_checkpoint(**hparams)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _(1) - Inference: Callbacks_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* _Test with LightingModule_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lightning Trainer\n",
    "trainer = pl.Trainer(callbacks=[GNNBuilder()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run TestStep\n",
    "trainer.test(model=model, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* _Test with LightningDataModule_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from Predict import SttDataModule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare LightningDataModule\n",
    "# dm = SttDataModule(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dm.setup(stage='test')\n",
    "# test_dataloaders = dm.test_dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run TestStep with LightningDataModule\n",
    "# trainer.test(model=model, dataloaders=None, ckpt_path=None, verbose=True, datamodule=dm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _(2) - Inference: Manual_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from Predict import eval_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- _Get Test Dataset from LightningModuel_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run setup() for datasets\n",
    "# model.setup(stage=\"fit\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get testset or test_dataloader\n",
    "# testset = model.testset\n",
    "# test_dataloader = model.test_dataloader()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- _Run `eval_model()` on `test_dataloader()`_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate model, returns torch tensors\n",
    "# scores, truths = eval_model(model, test_dataloader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _(3) - Inference: BNNBuilder_\n",
    "\n",
    "_If **GNNBuilder** callback has been run during training, just load data from `gnn_processed/test` and extract `scores` and `y_pid ~ truth` and simply run the following metrics_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fetch all files\n",
    "inputdir = \"run_all/gnn_processed/pred\"\n",
    "gnn_files = sorted(glob.glob(os.path.join(inputdir, \"*\")))\n",
    "print(\"Number of Files: \", len(gnn_files))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- _Load all `truth` and `scores` from the `testset` from the `DNN` stage_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scoresl, truthsl = [], []\n",
    "\n",
    "for e in range(len(gnn_files)):\n",
    "    \n",
    "    # logging\n",
    "    if e !=0 and e%1000==0:\n",
    "        print(\"Processed Batches: \", e)\n",
    "    \n",
    "    gnn_data = torch.load(gnn_files[e], map_location=device)\n",
    "    \n",
    "    truth = gnn_data.y_pid\n",
    "    score = gnn_data.scores\n",
    "    score = score[:truth.size(0)]\n",
    "    \n",
    "    # append each batch\n",
    "    scoresl.append(score)\n",
    "    truthsl.append(truth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# concatenate all\n",
    "scores = torch.cat(scoresl)\n",
    "truths = torch.cat(truthsl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save as .npy files\n",
    "np.save(\"gnn_scores.npy\", scores.numpy())\n",
    "np.save(\"gnn_truths.npy\", truths.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _Test Dataset_"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- _Get Data from LightningModule_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Method 1: Directly Get Test Dataset\n",
    "# testset = model.testset\n",
    "\n",
    "# Get singel Batch\n",
    "# batch = testset[0]\n",
    "\n",
    "# OR, loop over\n",
    "# for index, batch in enumerate(testset):\n",
    "# for batch in testset:\n",
    "#    print(index, batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Method 2: Directly Get Test Dataloader\n",
    "# test_dataloader = model.test_dataloader()\n",
    "\n",
    "# Get singel Batch\n",
    "# batch = next(iter(test_dataloader))\n",
    "\n",
    "# OR, loop over\n",
    "# for batch_idx, batch in enumerate(test_dataloader):\n",
    "# for batch in test_dataloader:\n",
    "#    print(batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- _Get Data from Test Dataset_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test Dataset from GNNBuilder\n",
    "inputdir=\"run_all/gnn_processed/test\"\n",
    "all_events = sorted(glob.glob(os.path.join(inputdir, \"*\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded_events = []\n",
    "for e in tqdm(all_events):\n",
    "    loaded_events.append(torch.load(e, map_location=device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PyG DataLoader\n",
    "test_dataloader = DataLoader(loaded_events, batch_size=1, num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fetch One Batch\n",
    "sampled_data = next(iter(test_dataloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print One Batch\n",
    "print(sampled_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _Plot Test Event_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "testset = model.testset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_data = testset[0]\n",
    "r, phi, ir = example_data.x.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = r * np.cos(phi * np.pi), r * np.sin(phi * np.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6,6))\n",
    "plt.scatter(x, y, s=2)\n",
    "plt.title(\"Azimuthal View of Detector\", fontsize=24), plt.xlabel(\n",
    "    \"x\", fontsize=18\n",
    "), plt.ylabel(\"y\", fontsize=18)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "e = example_data.edge_index\n",
    "pid = example_data.pid\n",
    "true_edges = pid[e[0]] == pid[e[1]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6,6))\n",
    "# plt.plot(x[e[:, ~true_edges]], y[e[:, ~true_edges]], c=\"r\")\n",
    "plt.plot(x[e[:, true_edges]], y[e[:, true_edges]], c=\"k\")\n",
    "plt.scatter(x, y, s=5)\n",
    "plt.title(\"Azimuthal View of Detector\", fontsize=24), plt.xlabel(\n",
    "    \"x\", fontsize=18\n",
    "), plt.ylabel(\"y\", fontsize=18)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6,6))\n",
    "plt.plot(x[e[:, (~true_edges)][:, 0:-1:5]], y[e[:, (~true_edges)][:, 0:-1:5]], c=\"r\")\n",
    "plt.scatter(x, y, s=5)\n",
    "plt.title(\"Azimuthal View of Detector\", fontsize=24), plt.xlabel(\n",
    "    \"x\", fontsize=18\n",
    "), plt.ylabel(\"y\", fontsize=18)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### _TensorBoard Logger_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load TensorBoard notebook extension\n",
    "# %load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %tensorboard --logdir=run_all/lightning_models/lightning_checkpoints/DNNStudy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}