train-gainEst.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Network Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Includes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ],
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# mass includes\n",
    "import os, sys, warnings\n",
    "import ipdb\n",
    "import torch as t\n",
    "import torchvision as tv\n",
    "import torchnet as tnt\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "# add paths for all sub-folders\n",
    "paths = [root for root, dirs, files in os.walk('.')]\n",
    "for item in paths:\n",
    "    sys.path.append(item)\n",
    "\n",
    "from ipynb.fs.full.config import mainConf\n",
    "from ipynb.fs.full.monitor import Visualizer\n",
    "from ipynb.fs.full.network import r2rNet, gainEst\n",
    "from ipynb.fs.full.dataLoader import fivekNight, valSet\n",
    "from ipynb.fs.full.util import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ],
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# for debugging only\n",
    "%pdb off\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# choose GPU if available\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = '0'\n",
    "device = t.device('cuda' if t.cuda.is_available() else 'cpu')\n",
    "\n",
    "# define models\n",
    "opt = mainConf()\n",
    "converter = r2rNet().to(device)\n",
    "converter.load('./saves')\n",
    "converter.eval()\n",
    "gain_est_model = gainEst().to(device)\n",
    "\n",
    "# load pre-trained model if necessary\n",
    "if opt.save_root:\n",
    "    _ = gain_est_model.load(opt.save_root)\n",
    "\n",
    "# dataloader for training\n",
    "train_dataset = fivekNight(opt)\n",
    "train_loader = t.utils.data.DataLoader(train_dataset,\n",
    "                                       batch_size=opt.batch_size,\n",
    "                                       shuffle=True,\n",
    "                                       num_workers=opt.num_workers,\n",
    "                                       pin_memory=True)\n",
    "\n",
    "# dataloader for validation\n",
    "val_dataset = valSet(opt)\n",
    "val_loader = t.utils.data.DataLoader(val_dataset)\n",
    "\n",
    "# optimizer\n",
    "bce_loss = t.nn.BCEWithLogitsLoss()\n",
    "l2_loss = t.nn.MSELoss()\n",
    "gain_est_optim = t.optim.Adam(gain_est_model.parameters(), lr=opt.lr)\n",
    "\n",
    "# visualizer\n",
    "vis = Visualizer(env='deepSelfie(gainEst)', port=8686)\n",
    "gain_est_meter = tnt.meter.AverageValueMeter()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "def validate():\n",
    "    # set to evaluation mode\n",
    "    gain_est_model.eval()\n",
    "\n",
    "    mask_error = 0.0\n",
    "    amp_error = 0.0\n",
    "    for (_, thumb_img, struct_img, seg_mask, amp, _, _, _) in val_loader:\n",
    "        with t.no_grad():\n",
    "            # copy to device\n",
    "            thumb_img = thumb_img.to(device)\n",
    "            struct_img = struct_img.to(device)\n",
    "            seg_mask = seg_mask.to(device)\n",
    "            amp = amp.to(device)\n",
    "\n",
    "            # inference\n",
    "            pred_mask, pred_amp = gain_est_model(thumb_img, struct_img)\n",
    "\n",
    "            # compute mse\n",
    "            mask_error += t.mean(\n",
    "                t.abs(t.nn.functional.sigmoid(pred_mask) - seg_mask))\n",
    "            amp_error += t.mean(t.abs(pred_amp - amp / opt.amp_range[1]))\n",
    "    mask_error /= len(val_loader)\n",
    "    amp_error /= len(val_loader)\n",
    "\n",
    "    # set to training mode\n",
    "    gain_est_model.train(mode=True)\n",
    "\n",
    "    return mask_error, amp_error"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training entry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "for epoch in range(0, 2):\n",
    "    # reset meter and gradient\n",
    "    gain_est_meter.reset()\n",
    "    gain_est_optim.zero_grad()\n",
    "\n",
    "    for index, (syth_img, syth_mask) in tqdm(enumerate(train_loader),\n",
    "                                             desc='epoch %d' % epoch,\n",
    "                                             total=len(train_loader)):\n",
    "        # copy to device\n",
    "        syth_img = syth_img.to(device)\n",
    "        syth_mask = syth_mask.to(device)\n",
    "\n",
    "        # convert to training sample\n",
    "        thumb_img, struct_img, seg_mask, amp, _, _, _ = toRaw(\n",
    "            converter, syth_img, syth_mask, opt)\n",
    "\n",
    "        # inference\n",
    "        pred_mask, pred_amp = gain_est_model(thumb_img, struct_img)\n",
    "\n",
    "        # compute loss\n",
    "        gain_est_loss = bce_loss(pred_mask, seg_mask) + l2_loss(\n",
    "            pred_amp, amp / opt.amp_range[1])\n",
    "\n",
    "        # compute gradient\n",
    "        gain_est_loss.backward()\n",
    "\n",
    "        # update parameter and reset gradient\n",
    "        gain_est_optim.step()\n",
    "        gain_est_optim.zero_grad()\n",
    "\n",
    "        # add to loss meter for logging\n",
    "        gain_est_meter.add(gain_est_loss.item())\n",
    "\n",
    "        # show intermediate result\n",
    "        if (index + 1) % opt.plot_freq == 0:\n",
    "            vis.plot('loss (gain est)', gain_est_meter.value()[0])\n",
    "            gain_est_plot = t.cat(\n",
    "                [seg_mask, t.nn.functional.sigmoid(pred_mask)],\n",
    "                dim=-1)[0, 0, :, :]\n",
    "            vis.img('gain est mask gt/pred', gain_est_plot.cpu() * 255)\n",
    "\n",
    "        # save model\n",
    "        if (index + 1) % opt.save_freq == 0:\n",
    "            gain_est_model.save()\n",
    "            mask_error, amp_error = validate()\n",
    "            vis.log('epoch: %d, err(mask/amp): %.4f, %.4f' %\n",
    "                    (epoch, mask_error, amp_error))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.3"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}