train-rawProcess.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, rawProcess\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",
    "raw_process_model = rawProcess().to(device)\n",
    "\n",
    "# load pre-trained model if necessary\n",
    "if opt.save_root:\n",
    "    last_epoch = raw_process_model.load(opt.save_root)\n",
    "else:\n",
    "    last_epoch = 0\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",
    "img_loss = imgLoss(device=device)\n",
    "raw_process_optim = t.optim.Adam(raw_process_model.parameters(), lr=opt.lr)\n",
    "\n",
    "# visualizer\n",
    "vis = Visualizer(env='deepSelfie(rawProcess)', port=8686)\n",
    "raw_process_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",
    "    raw_process_model.eval()\n",
    "\n",
    "    isp_psnr = 0.0\n",
    "    fuse_psnr = 0.0\n",
    "    for (syth_img, _, _, _, amp, noisy_raw, sorted_mask, wb) in val_loader:\n",
    "        with t.no_grad():\n",
    "            # copy to device\n",
    "            syth_img = syth_img.to(device)\n",
    "            amp = amp.to(device)\n",
    "            noisy_raw = noisy_raw.to(device)\n",
    "            sorted_mask = sorted_mask.to(device)\n",
    "            wb = wb.to(device)\n",
    "\n",
    "            # pre-processing\n",
    "            amp_high, _ = t.max(amp, 1)\n",
    "            amp_low, _ = t.min(amp, 1)\n",
    "\n",
    "            # inference\n",
    "            pred_high, pred_low, pred_fused = raw_process_model(\n",
    "                noisy_raw, amp_high, amp_low, wb)\n",
    "\n",
    "            # compute mse\n",
    "            pred_masked = sorted_mask[:, 0, :, :].unsqueeze(\n",
    "                1) * pred_high + sorted_mask[:,\n",
    "                                             1, :, :].unsqueeze(1) * pred_low\n",
    "            isp_mse = t.nn.functional.mse_loss(pred_masked, syth_img)\n",
    "            isp_psnr += 10 * t.log10(1 / isp_mse)\n",
    "            fuse_mse = t.nn.functional.mse_loss(pred_fused, syth_img)\n",
    "            fuse_psnr += 10 * t.log10(1 / fuse_mse)\n",
    "    isp_psnr /= len(val_loader)\n",
    "    fuse_psnr /= len(val_loader)\n",
    "\n",
    "    # set to training mode\n",
    "    raw_process_model.train(mode=True)\n",
    "\n",
    "    return isp_psnr, fuse_psnr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training entry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# reset meter and gradient\n",
    "raw_process_meter.reset()\n",
    "raw_process_optim.zero_grad()\n",
    "\n",
    "for index, (syth_img, syth_mask) in tqdm(enumerate(train_loader),\n",
    "                                         desc='progress',\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",
    "    _, _, _, amp, noisy_raw, sorted_mask, wb = toRaw(converter, syth_img,\n",
    "                                                     syth_mask, opt)\n",
    "    amp_high, _ = t.max(amp, 1)\n",
    "    amp_low, _ = t.min(amp, 1)\n",
    "\n",
    "    # inference\n",
    "    pred_high, pred_low, pred_fused = raw_process_model(\n",
    "        noisy_raw, amp_high, amp_low, wb)\n",
    "\n",
    "    # compute loss\n",
    "    pred_masked = sorted_mask[:, 0, :, :].unsqueeze(\n",
    "        1) * pred_high + sorted_mask[:, 1, :, :].unsqueeze(1) * pred_low\n",
    "    raw_process_loss = img_loss(pred_masked, pred_fused, syth_img)\n",
    "\n",
    "    # compute gradient\n",
    "    raw_process_loss.backward()\n",
    "\n",
    "    # update parameter and reset gradient\n",
    "    raw_process_optim.step()\n",
    "    raw_process_optim.zero_grad()\n",
    "\n",
    "    # add to loss meter for logging\n",
    "    raw_process_meter.add(raw_process_loss.item())\n",
    "\n",
    "    # show intermediate result\n",
    "    if (index + 1) % opt.plot_freq == 0:\n",
    "        vis.plot('loss (raw process)', raw_process_meter.value()[0])\n",
    "        raw_process_plot = t.nn.functional.interpolate(\n",
    "            t.clamp(t.cat([syth_img, pred_high, pred_low, pred_fused], dim=-1),\n",
    "                    0.0, 1.0),\n",
    "            scale_factor=0.5)[0, :, :, :]\n",
    "        vis.img('raw process gt/hi/lo/fuse', raw_process_plot.cpu() * 255)\n",
    "\n",
    "    # save model\n",
    "    if (index + 1) % opt.save_freq == 0:\n",
    "        raw_process_model.save()\n",
    "        isp_psnr, fuse_psnr = validate()\n",
    "        vis.log('psnr(isp/fuse): %.2f, %.2f' % (isp_psnr, fuse_psnr))"
   ]
  }
 ],
 "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
}