train_main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eaaad3f4-0894-4877-801c-630f27e9e1e4",
   "metadata": {},
   "source": [
    "# Super resolution notebook : training part\n",
    "The slc data we want to use the super resolution on is stored in './data/train/slc_lr' and the denoised is stored in './data/train/denoised_lr'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f738135b-5270-426a-9487-a73cc466572d",
   "metadata": {},
   "source": [
    "## Parameters to modify\n",
    "- device : 'cuda' to use gpu, 'cpu' otherwise\n",
    "- model_name : network to use for super resolution, choose between :\n",
    "    - 'SWINTRANSFORMER'\n",
    "    - 'SRCNN'\n",
    "    - 'PixelShuffle'\n",
    "- SR_type : super resolution method, choose between :\n",
    "    - 'n_n' for noisy input to noisy output\n",
    "    - 'n_d' for noisy input to denoised output\n",
    "    - 'd_d' for denoised input to denoised output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0248ad65-5808-40c1-bedf-6b69171e5c62",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda' # cuda or cpu\n",
    "model_name = 'SWINTRANSFORMER' # SRCNN, PixelShuffle, SWINTRANSFORMER \n",
    "SR_type = 'n_d' # n_n, n_d or d_d "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53c47f9d-a636-4926-a01a-accde981ae18",
   "metadata": {},
   "source": [
    "## Run all these cells to make super resolution with the parameters chosed above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "50f80fc3-d143-465b-96b8-d4fd3856e1b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import yaml\n",
    "import argparse\n",
    "import shutil\n",
    "import pathlib\n",
    "import torch\n",
    "import sys\n",
    "import numpy as np\n",
    "from torchsummary import summary\n",
    "from matplotlib import pyplot as plt\n",
    "from utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cfa94a5f-d169-4ba5-a191-120324f01381",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/cfren/anaconda3/envs/recherche/lib/python3.11/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:3526.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SWINTRANSFORMER network loaded\n",
      "1 train image(s) loaded\n",
      "0 valid image(s) loaded\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/cfren/anaconda3/envs/recherche/lib/python3.11/site-packages/torch/utils/data/dataset.py:414: UserWarning: Length of split at index 1 is 0. This might result in an empty dataset.\n",
      "  warnings.warn(f\"Length of split at index {i} is 0. \"\n"
     ]
    }
   ],
   "source": [
    "# load config\n",
    "cfg = None\n",
    "valid_flag = True\n",
    "\n",
    "with open('./config_base.yaml', 'r') as y_file:\n",
    "    cfg = yaml.load(y_file, Loader=yaml.FullLoader)\n",
    "\n",
    "# load model\n",
    "\n",
    "model = get_model(cfg,model_name,device)\n",
    "model = model.to(device)\n",
    "\n",
    "if cfg[\"TRAIN\"][\"SUMMARY\"]:\n",
    "    print(summary(\n",
    "        model,\n",
    "        input_size=(\n",
    "            cfg[\"DATASET\"][\"IN_CHANNELS\"],\n",
    "            cfg[\"DATASET\"][\"IMAGE_SIZE\"] // cfg[\"DATASET\"][\"PREPROCESSING\"]['DOWNSCALE_FACTOR'],\n",
    "            cfg[\"DATASET\"][\"IMAGE_SIZE\"] // cfg[\"DATASET\"][\"PREPROCESSING\"]['DOWNSCALE_FACTOR'],\n",
    "        ),\n",
    "    ))\n",
    "\n",
    "if cfg[\"TRAIN\"][\"PATH_TO_WEIGHTS\"] != 'None':\n",
    "    load_path = os.path.join(cfg[\"TRAIN\"][\"PATH_TO_WEIGHTS\"],'{}_{}.pth'.format(model_name.lower(),SR_type))\n",
    "    load_network(model, load_path, strict=True)\n",
    "    \n",
    "# save dir\n",
    "save_path = os.path.join(cfg[\"TRAIN\"][\"PATH_TO_SAVE\"],'{}_{}'.format(model_name.lower(),SR_type))\n",
    "os.makedirs(save_path,exist_ok=True)\n",
    "\n",
    "\n",
    "im_path,p_lr,p_hr = get_path_param(SR_type,cfg)\n",
    "train_data, valid_data = load_train(cfg, p_lr)\n",
    "\n",
    "print('{} network loaded'.format(model_name))\n",
    "print('{} train image(s) loaded'.format(len(train_data)))\n",
    "print('{} valid image(s) loaded'.format(len(valid_data)))\n",
    "\n",
    "if len(valid_data) == 0:\n",
    "    #For dummy training case\n",
    "    valid_flag = False\n",
    "\n",
    "# Define the loss\n",
    "f_loss = get_loss(cfg).to(device)\n",
    "\n",
    "# Define the optimizer\n",
    "optimizer = get_optimizer(cfg, model.parameters())\n",
    "\n",
    "# Define Scheduler\n",
    "scheduler = get_scheduler(cfg, optimizer)\n",
    "\n",
    "# Init Checkpoint class\n",
    "checkpoint = ModelCheckpoint(\n",
    "    save_path, model, cfg[\"TRAIN\"][\"EPOCH\"], cfg[\"TRAIN\"][\"CHECKPOINT_STEP\"]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8ccad76e-fb35-45fa-b37a-3f4013cb5d58",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  1.59it/s]\n",
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  6.23it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH : 0, Train loss: 0.07482893764972687, Valid loss: 0.035746168345212936\n",
      "PSNR: 59.619692602828934, L1: 0.20908474922180176, L2: 0.0715336948633194, SSIM: 0.6776777505874634, HUBER loss: 0.035746168345212936\n",
      "Saving a better model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  7.86it/s]\n",
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  4.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH : 1, Train loss: 0.035469163209199905, Valid loss: 0.01827661134302616\n",
      "PSNR: 62.531589796633476, L1: 0.14518696069717407, L2: 0.03658650815486908, SSIM: 0.6613328456878662, HUBER loss: 0.01827661134302616\n",
      "Saving a better model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  7.32it/s]\n",
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  5.02it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH : 2, Train loss: 0.017962371930480003, Valid loss: 0.017339743673801422\n",
      "PSNR: 62.759037109434544, L1: 0.13680922985076904, L2: 0.03471972048282623, SSIM: 0.6323990821838379, HUBER loss: 0.017339743673801422\n",
      "Saving a better model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  7.28it/s]\n",
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  5.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH : 3, Train loss: 0.01754223182797432, Valid loss: 0.017736375331878662\n",
      "PSNR: 62.65902537986343, L1: 0.13929827511310577, L2: 0.03552854061126709, SSIM: 0.6311630010604858, HUBER loss: 0.017736375331878662\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  6.73it/s]\n",
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  5.42it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH : 4, Train loss: 0.018412185832858086, Valid loss: 0.01586924120783806\n",
      "PSNR: 63.13939617234185, L1: 0.1310785412788391, L2: 0.031808286905288696, SSIM: 0.6515644788742065, HUBER loss: 0.01586924120783806\n",
      "Saving a better model\n",
      "Saving model at Epoch 4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  7.83it/s]\n",
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  5.47it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH : 5, Train loss: 0.016288571059703827, Valid loss: 0.014085243456065655\n",
      "PSNR: 63.6558429324483, L1: 0.12035776674747467, L2: 0.028242014348506927, SSIM: 0.6728674173355103, HUBER loss: 0.014085243456065655\n",
      "Saving a better model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  7.61it/s]\n",
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  5.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH : 6, Train loss: 0.014278806746006012, Valid loss: 0.01310816965997219\n",
      "PSNR: 63.96836978792152, L1: 0.1153903678059578, L2: 0.02628106251358986, SSIM: 0.6799975633621216, HUBER loss: 0.01310816965997219\n",
      "Saving a better model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  7.56it/s]\n",
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  5.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH : 7, Train loss: 0.013055453076958656, Valid loss: 0.01297856867313385\n",
      "PSNR: 64.01332748388941, L1: 0.11699434369802475, L2: 0.026010408997535706, SSIM: 0.6782691478729248, HUBER loss: 0.01297856867313385\n",
      "Saving a better model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  7.81it/s]\n",
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  5.22it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH : 8, Train loss: 0.012936566025018692, Valid loss: 0.01305672898888588\n",
      "PSNR: 63.9893511404087, L1: 0.11968369781970978, L2: 0.026154402643442154, SSIM: 0.679003119468689, HUBER loss: 0.01305672898888588\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  7.07it/s]\n",
      "100%|█████████████████████████████████████████████| 1/1 [00:00<00:00,  5.34it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH : 9, Train loss: 0.013091996312141418, Valid loss: 0.012797551229596138\n",
      "PSNR: 64.07801877056299, L1: 0.11972814798355103, L2: 0.025625836104154587, SSIM: 0.684995174407959, HUBER loss: 0.012797551229596138\n",
      "Saving a better model\n",
      "Saving model at Epoch 9\n"
     ]
    }
   ],
   "source": [
    "# Start training loop\n",
    "for epoch in range(cfg[\"TRAIN\"][\"EPOCH\"]):\n",
    "    \n",
    "    # Train\n",
    "    training_loss = train_one_epoch(\n",
    "        model,\n",
    "        train_data,\n",
    "        f_loss,\n",
    "        optimizer,\n",
    "        device,\n",
    "    )\n",
    "\n",
    "    # Validation\n",
    "    if valid_flag:\n",
    "        (valid_loss, psnr, _, _, _, l1_loss, l2_loss, ssim_loss, huber_loss,) = valid_one_epoch(\n",
    "            model,\n",
    "            valid_loader,\n",
    "            f_loss,\n",
    "            device,\n",
    "            cfg[\"DATASET\"][\"CLIP\"][\"MAX\"] - cfg[\"DATASET\"][\"CLIP\"][\"MIN\"],\n",
    "        )\n",
    "        print(f\"EPOCH : {epoch}, Train loss: {training_loss}, Valid loss: {valid_loss}\")\n",
    "        print(f\"PSNR: {psnr}, L1: {l1_loss}, L2: {l2_loss}, SSIM: {ssim_loss}, HUBER loss: {huber_loss}\")\n",
    "    else:\n",
    "        #Dummy validation when no valid set\n",
    "        (valid_loss, psnr, _, _, _, l1_loss, l2_loss, ssim_loss, huber_loss,) = valid_one_epoch(\n",
    "            model,\n",
    "            train_data,\n",
    "            f_loss,\n",
    "            device,\n",
    "        )\n",
    "        print(f\"EPOCH : {epoch}, Train loss: {training_loss}, Valid loss: {valid_loss}\")\n",
    "        print(f\"PSNR: {psnr}, L1: {l1_loss}, L2: {l2_loss}, SSIM: {ssim_loss}, HUBER loss: {huber_loss}\")\n",
    "        # Update scheduler\n",
    "    if scheduler is not None:\n",
    "        scheduler.step(valid_loss)\n",
    "\n",
    "    # Save best model\n",
    "    checkpoint.update(valid_loss, epoch)\n",
    "\n",
    "    # Get current learning rate\n",
    "    learning_rate = optimizer.param_groups[0][\"lr\"]"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}