Tutorial/Noob Heart (Pytorch).ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NoobHeart's Pytorch implementation  \n",
    "\n",
    "PhysBench does not rely on any deep learning framework, it handles all preprocessing, postprocessing, visualization, and then exposes the core model development completely to users who can flexibly choose their own development tools.  \n",
    "\n",
    "### Preparation  \n",
    "\n",
    "First, you need to prepare the dataset and training data. Please refer to `Noob Heart.ipynb`, this part is the same.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"..\")\n",
    "from utils import * \n",
    "\n",
    "# To resolve the conflict between torch and numpy, sometimes it can cause a core dump.\n",
    "os.environ[\"KMP_DUPLICATE_LIB_OK\"] = \"TRUE\"\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from torch import optim"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize the convolutional kernel weights of NoobHeart and define the model structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NoobHeart(nn.Module):\n",
    "\n",
    "    def __init__(self,):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv3d(3, 4, (2, 2, 2), stride=(1, 2, 2))\n",
    "        self.conv2 = nn.Conv3d(4, 2, (2, 2, 2), stride=(1, 2, 2), padding=[1,0,0,])\n",
    "        self.conv3 = nn.Conv3d(2, 1, (1, 1, 1), stride=(1, 1, 1))\n",
    "        self.pool  = nn.AvgPool3d((1, 2, 2))\n",
    "        self.N = lambda x, ax=2:(x-torch.unsqueeze(x.mean(axis=ax), ax))/torch.unsqueeze(x.std(axis=ax), ax)\n",
    "\n",
    "    def __call__(self, x):\n",
    "        x = self.N(x.permute(0, 4, 1, 2, 3))\n",
    "        x = self.N(torch.tanh(self.conv1(x)))\n",
    "        x = self.N(torch.tanh(self.conv2(x)))\n",
    "        x = self.pool(x)\n",
    "        x = self.conv3(x)\n",
    "        return x.view(x.shape[0], -1)\n",
    "    \n",
    "model = NoobHeart()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write the training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch=0,\tval_loss=0.6574487686157227\n",
      "Best model saved\n",
      "epoch=1,\tval_loss=0.5800231695175171\n",
      "Best model saved\n",
      "epoch=2,\tval_loss=0.5603955984115601\n",
      "Best model saved\n",
      "epoch=3,\tval_loss=0.5514898300170898\n",
      "Best model saved\n",
      "epoch=4,\tval_loss=0.5484949350357056\n",
      "Best model saved\n",
      "epoch=5,\tval_loss=0.5488275289535522\n",
      "epoch=6,\tval_loss=0.5489773750305176\n",
      "epoch=7,\tval_loss=0.5491269826889038\n",
      "epoch=8,\tval_loss=0.5497186183929443\n",
      "epoch=9,\tval_loss=0.5492029190063477\n"
     ]
    }
   ],
   "source": [
    "batch = 32\n",
    "train = load_datatape(\"train_tape.h5\", batch=batch) # Training set\n",
    "valid = load_datatape(\"valid_tape.h5\", batch=batch) # Validation set\n",
    "\n",
    "optimizer = optim.Adam(model.parameters())\n",
    "\n",
    "model.cuda()\n",
    "best = None\n",
    "for epoch in range(10): # 10 epochs\n",
    "    model.train()\n",
    "    for step, (data, label) in enumerate(train): # train\n",
    "        data, label = torch.tensor(data).float().cuda(), torch.tensor(label).float().cuda()\n",
    "        optimizer.zero_grad()\n",
    "        loss = (model(data)-label).abs().mean()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    model.eval()\n",
    "    vloss = []\n",
    "    for step, (data, label) in enumerate(valid): # valid\n",
    "        data, label = torch.tensor(data).float().cuda(), torch.tensor(label).float().cuda()\n",
    "        vloss.append((model(data)-label).abs().mean())\n",
    "    vloss = torch.stack(vloss).mean()\n",
    "    print(f'{epoch=},\\tval_loss={vloss}')\n",
    "    if best is None or best>vloss:\n",
    "        torch.save(model.state_dict(), 'NoobHeart.pkl')\n",
    "        print('Best model saved')\n",
    "        best = vloss\n",
    "\n",
    "model.load_state_dict(torch.load('NoobHeart.pkl')) # load the best\n",
    "model_ = lambda x:model(torch.tensor(x.astype(np.float32)).cuda()).detach().cpu().numpy()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use `eval_on_dataset` to test and obtain metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2701bc20a544493c967cf025d9d42b18",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/42 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HR metrics: MAE:0.977, RMSE:1.477, R:0.99672\n",
      "HRV metrics: MAE:35.889, RMSE:40.899, R:0.63317\n"
     ]
    }
   ],
   "source": [
    "eval_on_dataset('ubfc_dataset.h5', model_, 32, (8, 8), step=1, batch=32, save='../results/NoobHeart_PURE_UBFC.h5')\n",
    "r = get_metrics('../results/NoobHeart_PURE_UBFC.h5')['Whole video']\n",
    "print(f'HR metrics: MAE:{r[\"MAE\"]}, RMSE:{r[\"RMSE\"]}, R:{r[\"R\"]}')\n",
    "r = get_metrics_HRV('../results/NoobHeart_PURE_UBFC.h5')['SDNN']\n",
    "print(f'HRV metrics: MAE:{r[\"MAE\"]}, RMSE:{r[\"RMSE\"]}, R:{r[\"R\"]}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dl",
   "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.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}