ResCSNet_5_2_1_github.ipynb

{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"name":"ResCSNet_5_2_1_github.ipynb","version":"0.3.2","provenance":[],"collapsed_sections":[],"toc_visible":true},"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.6.8"},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"accelerator":"GPU"},"cells":[{"cell_type":"markdown","metadata":{"colab_type":"text","id":"Y2tiMNrcZx5h"},"source":["# ResCSNet\n","Inspired by ConvCSNet. Can I just port the ResNet Multibranch structure in?"]},{"cell_type":"markdown","metadata":{"id":"qdDPNkwSp-Vv","colab_type":"text"},"source":["# About this notebook\n","This notebook is intended to run in [Google Colaboratory](https://colab.research.google.com). It may require a lot changes (mostly deletions) if you want to run it from your local device.\n","\n","In order to view tensorboard plots in the Colab VM during trainning, I have applied some dirty hacks (using `frpc` and a remote VPS running `frpd`). Also I am using `pydrive` module to download dataset from my Google Drive, and upload the model checkpoint at the end of each epoch.\n","\n","**Update**: Tensorboard 2.0 has added a \"inline\" tensorboard magic for jupyter notebooks. It is recommended that you open **another** notebook which shares the same VM with this notebook and run the following:\n","```\n","!pip install -q tf-nightly-2.0-preview\n","# Load the TensorBoard notebook extension\n","%load_ext tensorboard\n","%tensorboard --logdir runs\n","```\n","In this way you don't have to bother with a VPS or `frpc` or something."]},{"cell_type":"markdown","metadata":{"id":"1je7BvI-p-Vy","colab_type":"text"},"source":["# About the dataset\n","I will load the pictures from the COCO dataset downloaded (and grayscaled and center-cropped already) by myself. You may download it from https://drive.google.com/open?id=12Nje-yhxcIVyz7L_lVxxcfXcTWa5Ba-m\n","\n","Some of the code below may try to download it from (your) Google Drive. It may be better to upload the file to your Google Drive."]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"v8MZcUaDXR4s"},"source":["# Install necessary packages and download dataset"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"cmOC2bSLljL-","colab":{}},"source":["!pip install -U -q tensorboardX\n","!pip install -U -q PyDrive\n","\n","from pydrive.auth import GoogleAuth\n","from pydrive.drive import GoogleDrive\n","from google.colab import auth\n","from oauth2client.client import GoogleCredentials\n","\n","# Authenticate and create the PyDrive client.\n","# This only needs to be done once per notebook.\n","auth.authenticate_user()\n","gauth = GoogleAuth()\n","gauth.credentials = GoogleCredentials.get_application_default()\n","drive = GoogleDrive(gauth)\n","\n","# List .gz files in the root.\n","#\n","# Search query reference:\n","# https://developers.google.com/drive/v2/web/search-parameters\n","listed = drive.ListFile({'q': \"title contains '.gz' and 'root' in parents\"}).GetList()\n","for file in listed:\n","  print('title {}, id {}'.format(file['title'], file['id']))\n","\n","downloaded = drive.CreateFile({'id': \"12Nje-yhxcIVyz7L_lVxxcfXcTWa5Ba-m\"})\n","downloaded.GetContentFile(\"center-crop-100.tar.gz\")\n","\n","# Unextract dataset\n","print(\"Extract dataset\")\n","!tar -xzf \"center-crop-100.tar.gz\"\n","\n","\n","from pydrive.auth import GoogleAuth\n","from pydrive.drive import GoogleDrive\n","from google.colab import auth\n","from oauth2client.client import GoogleCredentials\n","\n","# Authenticate and create the PyDrive client.\n","print(\"Authenticate and create the PyDrive client\")\n","auth.authenticate_user()\n","_gauth = GoogleAuth()\n","_gauth.credentials = GoogleCredentials.get_application_default()\n","# drive = GoogleDrive(_gauth)\n","\n","# Create a file instance to upload checkpoint later\n","# gfile_ckpt = drive.CreateFile()\n","\n","# pytorch-wavelets-1.0.0 is known to be OK. Higher versions may work as well\n","!git clone \"https://github.com/fbcotter/pytorch_wavelets.git\"\n","!pip install ./pytorch_wavelets\n","print(\"Done\")"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"LloheIUNXbux"},"source":["# Download exsiting parameters\n","Needed if you want to resume trainning or test with exsiting parameters.\n","\n","Some pretrained models:\n","[ResCsNet-colab-5_2_1-r0.25_checkpoint.pth](https://drive.google.com/open?id=1QQJJ3c9SlMK03v_v28J_K0ymvvA2vbhG), [ResCsNet-colab-5_2_1-r0.20_checkpoint.pth](https://drive.google.com/open?id=12g8reeeyi4Ei8v9dZudS_jsRE8ImYuq5), [ResCsNet-colab-5_2_1-r0.15_checkpoint.pth](https://drive.google.com/open?id=1J2tOS02BJVBWwOT-p2SDwpKRuociEKPO), [ResCsNet-colab-5_2_1-r0.10_checkpoint.pth](https://drive.google.com/open?id=1oLLB45GWfGxhdxYqNoYPQRcNX5_v9XFE), [ResCsNet-colab-5_2_1-r0.05_checkpoint.pth](https://drive.google.com/open?id=1xJbTgKgjqUwyQJn8gxvWy7rJ5N-9maWu)"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"Eb5hcAVKnNn7","colab":{}},"source":["drive = GoogleDrive(gauth)\n","# download trained params\n","listed = drive.ListFile({'q': \"title contains 'ResCsNet-colab-5_2_1-r0.25_checkpoint.pth' and trashed=False\"}).GetList()\n","for file in listed:\n","  print('title {}, id {}'.format(file['title'], file['id']))\n","# listed = drive.ListFile({'q': \"title contains 'ResCsNet-colab-5_2_1-r0.20_checkpoint.pth' and trashed=False\"}).GetList()\n","# for file in listed:\n","#   print('title {}, id {}'.format(file['title'], file['id']))"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"faTyL8u2oBoH","colab":{}},"source":["drive = GoogleDrive(gauth)\n","downloaded2 = drive.CreateFile({'id': \"1QQJJ3c9SlMK03v_v28J_K0ymvvA2vbhG\"})\n","downloaded2.GetContentFile(\"ResCsNet-colab-5_2_1-r0.25_checkpoint.pth\")\n","# downloaded3 = drive.CreateFile({'id': \"12g8reeeyi4Ei8v9dZudS_jsRE8ImYuq5\"})\n","# downloaded3.GetContentFile(\"ResCsNet-colab-5_2_1-r0.20_checkpoint.pth\")"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"5IxrvRoYmxcH"},"source":["-----------------------------------------------------------------"]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"b-AycMcxXkbC"},"source":["# Runtime configurations\n","This section mostly covers hacks and tricks. If you are using tensorflow 2.0 and inline tensorboard you probably do not need to run cells in this section."]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"F1T9ej3WXqph"},"source":["## nvidia-smi"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"CrdJak_OgiPx","colab":{}},"source":["!nvidia-smi"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"LghacfEZXpEE"},"source":["## tensorboard"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"S6ZpjvMwBXUD","colab":{}},"source":["!tar -xjvf runs.tbz"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"BvK4_rdmLnIU","colab":{}},"source":["get_ipython().system_raw(\"tensorboard --logdir runs --host 127.0.0.1 &\")\n","!ps -ef | grep tensorboard"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"tRHqCyCOf2BF"},"source":["## python http.server"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"T_mlYw-vf4v_","colab":{}},"source":["get_ipython().system_raw(\"python3 -m http.server 8000 --bind 127.0.0.1 &\")\n","!ps -ef | grep http"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"mZS5gRkUXtYi"},"source":["## frpc"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"FAmkMgkY6uQl","colab":{}},"source":["!wget \"https://github.com/fatedier/frp/releases/download/v0.24.1/frp_0.24.1_linux_amd64.tar.gz\"\n","!mkdir frp\n","!tar -xvf \"frp_0.24.1_linux_amd64.tar.gz\" -C frp"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"NohC-gdf7Rsk","colab":{}},"source":["get_ipython().system_raw(\"./frp/frp_0.24.1_linux_amd64/frpc -c ./frpc.ini &\")\n","#!./frp/frp_0.24.1_linux_amd64/frpc -c ./frpc.ini"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"AmD9duEV9E_m","colab":{}},"source":["!ps -ef | grep frpc"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"Ehk9Q9ilAyeU","colab":{}},"source":["!cat frpc.ini"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"WRCUVgQzXwiL"},"source":["## Control the go and stop of the trainning"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"yqe4ohRvBZXz","colab":{}},"source":["# tell the program to stop trainning\n","# !touch _stop\n","\n","# or lift the ban\n","# !rm _stop"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"72F4eitcX2gi"},"source":["## tensorboard data archive"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"OZ9BRjIG4PYi","colab":{}},"source":["get_ipython().system_raw(\"bash get_runs.sh &\")\n","# get_ipython().system_raw(\"python3 upload_runs.py &\")\n","#!bash get_runs.sh\n","#!python3 upload_runs.py"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"yAx5a55w51NT","colab":{}},"source":["!ps -ef | grep get_runs\n","# !ps -ef | grep upload_runs.py"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"HpwZOj8Sdod3"},"source":["## tail trainning log"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"sQzoZhHhdquc","colab":{}},"source":["# !tail \"ResCsNet-colab-5_2_1-r0.10_trainning.log\""],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"gBu8dfIbX7MC"},"source":["## Google drive re-auth\n","I am doing this because sometimes I encouter bugs if not re-auth with Google drive. Not sure if I missed something in pydrive documention."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"S5cpqPSlCsNB","colab":{}},"source":["# if the upload has to fail I will do that manually\n","from pathlib import Path\n","from pydrive.auth import GoogleAuth\n","from pydrive.drive import GoogleDrive\n","from google.colab import auth\n","from oauth2client.client import GoogleCredentials\n","\n","# !mkdir authentication\n","\n","gauth = GoogleAuth()\n","gauth.LoadCredentialsFile(\"mycreds.txt\")\n","if gauth.credentials is None:\n","    # Authenticate if they're not there\n","    # self.gauth.LocalWebserverAuth()\n","    print(\"no creds saved\")\n","    auth.authenticate_user()\n","    gauth.credentials = GoogleCredentials.get_application_default()\n","elif gauth.access_token_expired:\n","    # Refresh them if expired\n","    print(\"token expired\")\n","    gauth.Refresh()\n","else:\n","    # Initialize the saved creds\n","    print(\"Initialize the saved creds\")\n","    gauth.Authorize()\n","# Save the current credentials to a file\n","gauth.SaveCredentialsFile(\"mycreds.txt\") \n","\n","\n","############################################\n","# the_drive = GoogleDrive(gauth)\n","# ckpt_name = \"ResCsNet-colab-5_2_1-r0.10_checkpoint.pth\"\n","\n","\n","# id_file = Path('_id')\n","# if id_file.exists():\n","#     print(\"_id exsits\")\n","#     fileid = id_file.read_text()\n","#     _gfile_ckpt = the_drive.CreateFile({'id': fileid})\n","#     _gfile_ckpt.SetContentFile(ckpt_name)\n","#     _gfile_ckpt.Upload()\n","# else:\n","#     print(\"_id not exsits\")\n","#     _gfile_ckpt = the_drive.CreateFile()      \n","#     _gfile_ckpt.SetContentFile(ckpt_name)\n","#     _gfile_ckpt.Upload()\n","#     id_file.write_text(_gfile_ckpt['id']) \n","print(\"Done\")"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"kOKcWKOHha9Q"},"source":["## Check uptime"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"jAilc2gChdil","colab":{}},"source":["!uptime"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"uATqIq5e1iZt"},"source":["# ==========================="]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"jS4s9P9wcuCX"},"source":["# The real things"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"GIDa9XAfZx6R","scrolled":true,"colab":{}},"source":["# imports\n","from pathlib import Path\n","from PIL import Image, ImageFile\n","# see https://stackoverflow.com/questions/12984426/python-pil-ioerror-image-file-truncated-with-big-images\n","ImageFile.LOAD_TRUNCATED_IMAGES = True\n","\n","# from six.moves import cPickle as pickle\n","# import platform\n","\n","import torch\n","import torch.nn as nn\n","import torch.optim as optim\n","from torch.utils.data import Dataset, DataLoader, random_split, Subset\n","from torch.utils.data import sampler\n","import torch.nn.functional as F\n","\n","# import torchvision.datasets as dset\n","import torchvision.transforms as T\n","\n","import numpy as np\n","import matplotlib.pyplot as plt\n","from tensorboardX import SummaryWriter \n","\n","# from trainning_func import get_evaluation\n","\n","\n","# import ipdb\n","\n","%matplotlib inline\n","# %matplotlib tk\n","# plt.rcParams['figure.figsize'] = (10.0, 8.0) # set default size of plots\n","# plt.rcParams['image.interpolation'] = 'nearest'\n","# plt.rcParams['image.cmap'] = 'gray'\n","\n","# for auto-reloading external modules\n","# see http://stackoverflow.com/questions/1907993/autoreload-of-modules-in-ipython\n","%load_ext autoreload\n","%autoreload 2\n","\n","print(\"Done\")"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"rOocHDs3Zx6d"},"source":["## Dataset classes"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"UDP3J7LrZx6h","scrolled":true,"colab":{}},"source":["input_width_height = 100\n","# build somthing like an image pyramid. For details see class definition of ResCsNet\n","TOTAL_PIXELS = input_width_height*input_width_height   # N\n","SAMPLING_RATE = 0.25\n","SAMPLED_PIXELS = int(TOTAL_PIXELS * SAMPLING_RATE)     # M\n","WIDTH_INITIAL = int(np.ceil(np.sqrt(SAMPLED_PIXELS)))\n","WIDTH_HALF = int(input_width_height/4*2)\n","WIDTH_THREE_QUARTERS = int(input_width_height/4*3)\n","\n","class CocoDataset(Dataset):\n","    def __init__(self, path_dataset, transform=None):\n","        p_iter = Path(path_dataset).iterdir()\n","        self.transform = transform\n","        self.images_list = [name for name in p_iter]\n","    \n","    def __len__(self):\n","        '''provides the size of the dataset'''\n","        return len(self.images_list)\n","    \n","    def __getitem__(self, idx):\n","        '''\n","        supporting integer indexing in range from 0 to len(self) exclusive\n","        '''\n","        im = Image.open(self.images_list[idx]).convert(\"L\")  # to grayscale\n","        if self.transform:\n","            im = self.transform(im)\n","        return im"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"cTYC6f_VZx64"},"source":["## Dataset instances"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"dq1L7KBkZx67","scrolled":true,"colab":{}},"source":["# Test if CocoDataset is OK\n","# path_dset_train = '/home/xzc/center-crop-100/train2014'\n","# path_dset_test = '/home/xzc/center-crop-100/val2014'\n","path_dset_train = 'center-crop-100/train2014'\n","path_dset_test = 'center-crop-100/val2014'\n","# path_dset_train = r'D:\\dev_workspace\\CS-DL\\center-crop-100\\train2014'\n","# path_dset_test = r'D:\\dev_workspace\\CS-DL\\center-crop-100\\val2014'\n","input_width_height = 100\n","\n","train_set = CocoDataset(path_dset_train,\n","                         transform=T.Compose(\n","                            [\n","                                T.RandomHorizontalFlip(),\n","                                T.RandomVerticalFlip(),\n","                                T.ToTensor(),\n","                                T.Normalize((0.5,), (0.5,))\n","                            ])\n","                        )\n","test_set = CocoDataset(path_dset_test,\n","                       transform=T.Compose(\n","                            [\n","                                T.ToTensor(),\n","                                T.Normalize((0.5,), (0.5,))\n","                            ])\n","                      )\n","\n","print(\"Done\")"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"EdPjR37KZx6_","scrolled":false,"colab":{}},"source":["# function to denomalize normalized image\n","def denormalize(im, mean, std):\n","    '''\n","    im: pytorch tensor view as image\n","    '''\n","    assert len(im.size()) == 3\n","    if im.shape[0] == 1: # grayscale\n","#         im = im.reshape(im.shape[0], im.shape[1])\n","        im = im * std[0] + mean[0]\n","    else: # rgb\n","        for ch in range(3):\n","            im[:,:,ch] = im[:,:,ch] * std[ch] + mean[ch]\n","    return im\n","\n","def test_dataset():\n","    print(len(train_set))\n","    im = train_set[5]\n","\n","    print(im.shape)\n","    im = denormalize(im, (0.5,), (0.5,))\n","    print(im.shape)\n","\n","    _im = im.reshape((im.shape[1], im.shape[2]))\n","    #     plt.figure()\n","    plt.imshow(_im, cmap='gray')\n","    plt.show()\n","\n","# test_dataset()\n","print(\"Done\")"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"kZwxt7khZx7E","scrolled":true,"colab":{}},"source":["print(f\"{len(train_set)}, {len(test_set)}\")"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"K6tifxMnZx7N"},"source":["## Dataloader instances"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"C0-PMODhZx7P","scrolled":true,"colab":{}},"source":["# make the pytorch loader\n","\n","# final run: use this set\n","TOTAL_SAMPLES = len(train_set)\n","NUM_TRAIN = TOTAL_SAMPLES // 5 * 4\n","BATCH_SIZE = 64\n","\n","# mini test: use this set\n","# TOTAL_SAMPLES = 2000\n","# NUM_TRAIN = TOTAL_SAMPLES // 5 * 4\n","# BATCH_SIZE = 60\n","\n","# debug only\n","# TOTAL_SAMPLES = 100\n","# NUM_TRAIN = TOTAL_SAMPLES // 5 * 4\n","# BATCH_SIZE = 4\n","\n","loader_train = DataLoader(train_set, batch_size=BATCH_SIZE, sampler=sampler.SubsetRandomSampler(range(NUM_TRAIN)))\n","loader_subtrain = DataLoader(Subset(train_set, [i for i in range(3)]), batch_size=BATCH_SIZE)  # used for checking avg_psnr\n","loader_val = DataLoader(train_set, batch_size=BATCH_SIZE, sampler=sampler.SubsetRandomSampler(range(NUM_TRAIN, TOTAL_SAMPLES)))\n","loader_test = DataLoader(test_set, batch_size=BATCH_SIZE)\n","\n","print(f'len(loader_train)=={len(loader_train)}, len(loader_subtrain)={len(loader_subtrain)}')\n","print(f'len(loader_val)=={len(loader_val)}, len(loader_test)=={len(loader_test)}')\n","print('Done')"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"sg6PBKLQZx7W","scrolled":true,"colab":{}},"source":["# Test the usage of dataloaders\n","def test_dataloader():\n","    train_iter = iter(loader_train)\n","    original_im = next(train_iter)\n","\n","    print(type(original_im))\n","    print(original_im.size())\n","\n","    print(\"------------------\")\n","    test_iter = iter(loader_test)\n","    im = next(test_iter)\n","    print(type(im))\n","    print(im.size())\n","    \n","# test_dataloader()"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"8JvMwLeUGU0K"},"source":["## Set up device"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"HAK47qc_Zx7b","scrolled":true,"colab":{}},"source":["# set up device\n","# will use cuda if available\n","device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n","print('using device:', device)"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"bbuaHVjFZx7f"},"source":["# Encoder"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"DCoHiWBhZx7g","scrolled":true,"colab":{}},"source":["# original linear encoder\n","class EncoderLinear(nn.Module):\n","    def _init_weights(self, m):\n","#         print(m)\n","        if type(m) == nn.Conv2d:\n","            nn.init.kaiming_normal_(m.weight.data)\n","            if m.bias is not None:\n","                nn.init.constant_(m.bias.data, 0)\n","    \n","    def __init__(self, N, M):\n","        super().__init__()\n","        self.linear = nn.Linear(N, M)\n","        self.linear.apply(self._init_weights)\n","    \n","    def forward(self, x):\n","        x_vec = x.view(x.shape[0], 1, 1, -1)\n","        out = self.linear(x_vec)\n","        return out"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"e5mavALJZx7r","scrolled":true,"colab":{}},"source":["# make instance and test the model\n","def test_encoder():\n","#     train_iter = iter(loader_train)\n","#     im = train_iter.next()\n","#     encoder = Encoder(1, kernel_size=15, stride=1, padding=0)\n","    print(f\"input_width_height={input_width_height}\")\n","    N = input_width_height*input_width_height\n","    M = int(N * 0.30)\n","    encoder = EncoderLinear(N, M)\n","    im = torch.randn(BATCH_SIZE,1,input_width_height,input_width_height)\n","    with torch.no_grad():\n","        y = encoder(im)\n","    print(y.size())\n","\n","# test_encoder()"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"45-qtAUfZx7y"},"source":["# The Recovery Network (Decoder)"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"7XqumJgBZx7z","scrolled":true,"colab":{}},"source":["class TheUpsample(nn.Module):\n","    '''\n","    Wrapper for torch.nn.functional.interpolate\n","    (why don't they write a ready-for-use 'nn.Interpolate'?)\n","    '''\n","    def __init__(self, size):\n","        super().__init__()\n","        self.size = size\n","    \n","    def forward(self, x):\n","        return nn.functional.interpolate(x, size=self.size, mode='nearest')\n","\n","class ResCsNet(nn.Module):\n","#     def _init_weights(self, m):\n","#         '''\n","#         used to init weights.\n","#         '''\n","#         print(m)\n","#         if type(m) == nn.Conv2d:\n","#             nn.init.kaiming_normal_(m.weight.data)\n","#             nn.init.constant_(m.bias.data, 0)\n","#         if type(m) == nn.Linear:\n","#             nn.init.kaiming_normal_(m.weight.data)\n","#             nn.init.constant_(m.bias.data, 0)\n","#         if type(m) == nn.ConvTranspose2d:\n","#             nn.init.kaiming_normal_(m.weight.data)\n","#             nn.init.constant_(m.bias.data, 0)\n","\n","    def __init__(self, N, M):\n","#     def __init__(self, encoder_out_ch, encoder_ksize, encoder_stride, encoder_padding):\n","        super().__init__()\n","        # encoder\n","#         self.encoder= Encoder(encoder_out_ch, encoder_ksize, encoder_stride, encoder_padding)\n","        self.encoder= EncoderLinear(N, M)\n","        \n","        # upsample\n","        self.initial_width = int(np.ceil(np.sqrt(M))) # should be identical to WIDTH_INITIAL\n","        self.upsample0 = TheUpsample((1, int(self.initial_width**2)) )\n","#         upsample1_width = int(input_width_height/4)\n","#         self.upsample1 = TheUpsample((upsample1_width, upsample1_width))\n","        upsample2_width = WIDTH_HALF\n","        self.upsample2 = TheUpsample((upsample2_width, upsample2_width))\n","        upsample3_width = WIDTH_THREE_QUARTERS\n","        self.upsample3 = TheUpsample((upsample3_width, upsample3_width))\n","        # the last upsample should up sample to the original size\n","        self.upsample4 = TheUpsample((input_width_height, input_width_height))\n","        \n","        # conv_scale\n","        self.conv_scale1 = nn.Conv2d(96, 1, kernel_size=1, stride=1, padding=0)\n","        self.conv_scale2 = nn.Conv2d(96, 1, kernel_size=1, stride=1, padding=0)\n","        self.conv_scale3 = nn.Conv2d(96, 1, kernel_size=1, stride=1, padding=0)\n","        self.conv_scale4 = nn.Conv2d(96, 1, kernel_size=1, stride=1, padding=0)\n","        \n","        # units\n","        # unit 1\n","        self.unit1 = nn.Sequential(\n","#             nn.Conv2d(encoder_out_ch, 96, kernel_size=3, stride=1, padding=1),\n","            nn.Conv2d(1, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96),\n","            nn.LeakyReLU()\n","\n","        )\n","\n","        # unit2\n","        self.unit2a_branch1 = nn.Sequential(\n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96),\n","        )\n","        \n","        self.unit2a_branch2 = nn.Sequential(\n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96),\n","            nn.LeakyReLU(),\n","            \n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96),\n","            nn.LeakyReLU(),\n","            \n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96)\n","        )\n","        \n","        self.unit2b_branch2 = nn.Sequential(\n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96),\n","            nn.LeakyReLU(),\n","            \n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96),\n","            nn.LeakyReLU(),\n","            \n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96)\n","        )\n","        \n","        \n","        # unit 3\n","        # branch1 is identical pass\n","        self.unit3_branch2 = nn.Sequential(\n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96),\n","            nn.LeakyReLU(),\n","            \n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96),\n","            nn.LeakyReLU(),\n","            \n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96)\n","        )\n","        \n","        # unit 4\n","        self.unit4_branch1 = nn.Sequential(\n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96),\n","        )\n","        \n","        self.unit4_branch2 = nn.Sequential(\n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96),\n","            nn.LeakyReLU(),\n","            \n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96),\n","            nn.LeakyReLU(),\n","            \n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96)\n","        )\n","        \n","        # unit 5\n","        # branch 1 is identical pass\n","        self.unit5_branch2 = nn.Sequential(\n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96),\n","            nn.LeakyReLU(),\n","            \n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96),\n","            nn.LeakyReLU(),\n","            \n","            nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1),\n","            nn.BatchNorm2d(96)\n","        )\n","    \n","    def forward(self, x):\n","        x = self.encoder(x)\n","        x_vec = self.upsample0(x)\n","        x = x_vec.view(x.shape[0], 1, self.initial_width, self.initial_width)  # 23x23, 32x32, 45x45, 55x55\n","        \n","        # unit 1\n","        x = self.unit1(x)\n","#         x = self.upsample1(x)\n","        \n","        # unit 2\n","        x2a_1 = self.unit2a_branch1(x)\n","        x2a_2 = self.unit2a_branch2(x)\n","        x = x2a_1 + x2a_2\n","        x2b_1 = x.clone()\n","        x2b_2 = self.unit2b_branch2(x)\n","        x = x2b_1 + x2b_2\n","        \n","        x = self.upsample2(x)  # remove this layer for r=0.3\n","        # unit 3\n","        x3_1 = x.clone()\n","        x3_2 = self.unit3_branch2(x)\n","        x = x3_1 + x3_2\n","        \n","        x = self.upsample3(x)\n","        # unit 4\n","        x4_1 = self.unit4_branch1(x)\n","        x4_2 = self.unit4_branch2(x)\n","        x = x4_1 + x4_2\n","        \n","        x = self.upsample4(x)\n","        # unit5\n","        x5_1 = x.clone()\n","        x5_2 = self.unit5_branch2(x)\n","        x = x5_1 + x5_2\n","        im4 = self.conv_scale4(x)\n","        \n","        return im4\n","\n","############################################\n","def ResCsNet_init_weights(m):\n","        '''\n","        used to init weights. This is just a clone of the method in the class above. Provided for convenience\n","        '''\n","#         print(m)\n","        if type(m) == nn.Conv2d:\n","            nn.init.kaiming_normal_(m.weight.data)\n","            if m.bias is not None:\n","                nn.init.constant_(m.bias.data, 0)\n","        if type(m) == nn.Linear:\n","            nn.init.kaiming_normal_(m.weight.data)\n","            nn.init.constant_(m.bias.data, 0)\n","        if type(m) == nn.ConvTranspose2d:\n","            nn.init.kaiming_normal_(m.weight.data)\n","            if m.bias is not None:\n","                nn.init.constant_(m.bias.data, 0)"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"L1onn4NHZx74","scrolled":true,"colab":{}},"source":["# Make instance and check\n","# %debug \n","def test_rescsnet():\n","    print(f'using device: {device}')\n","    # 8, kernel_size=11, stride=5, padding=0 for r=0.28\n","#     model = ResCsNet(encoder_out_ch=1, encoder_ksize=15, encoder_stride=1, encoder_padding=0)\n","    N = input_width_height*input_width_height\n","    M = int(N * 0.05)\n","    model = ResCsNet(N, M)\n","    model.apply(ResCsNet_init_weights)\n","#     train_iter = iter(loader_train)\n","#     im = train_iter.next()\n","    im = torch.randn(BATCH_SIZE,1,input_width_height,input_width_height).to(device=device)\n","    with torch.no_grad():\n","        model.train()\n","        model.to(device=device)\n","        recovered_im = model(im)\n","    print(recovered_im.size())\n","    print(\"------------------\")\n","    with torch.no_grad():\n","        model.eval()\n","        model.to(device=device)\n","        recovered_im = model(im)\n","    print(recovered_im.size())\n","    del model\n","    \n","# test_rescsnet()"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"u1e_YIcQZx8B","scrolled":true,"colab":{}},"source":["print(torch.cuda.memory_cached())\n","torch.cuda.empty_cache()\n","print(torch.cuda.memory_cached())"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"SPzmdnNeZx8L"},"source":["# Training"]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"xx1DIIm-ZG4Q"},"source":["## Experiment setup"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"z1VHbepBZx8U","scrolled":true,"colab":{}},"source":["exp_name = 'ResCsNet-colab'\n","\n","print(f\"input_width_height={input_width_height}\")\n","N = input_width_height*input_width_height\n","M = int(N * SAMPLING_RATE)\n","model = ResCsNet(N, M)\n","# model.apply(ResCsNet_init_weights)  # not needed if continue trainning or testing"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"pHBSbIjPZG4f"},"source":["## Load exsiting params\n","This section is needed if you want to test or continue trainning a model.\n","\n","If you want to train from scratch, this section should be skipped."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"mdIbTx4sZx8e","scrolled":true,"colab":{}},"source":["#######################################\n","# Load the parameters to continue trainning if desired\n","# see https://github.com/pytorch/examples/blob/d6b52110bae32cbefeea6d4ffbf8cede98ac16fc/imagenet/main.py#L175\n","#######################################\n","\n","want_load_params = True\n","if want_load_params:\n","    #######################################\n","    # You need to make this correct\n","    fname = 'ResCsNet-colab-5_2_1-r0.25_checkpoint.pth'\n","    #######################################\n","    checkpoint = torch.load(fname)\n","    tfx_steps = checkpoint['tfx_steps']\n","    print(f\"tfx_steps is {tfx_steps}\")\n","    tfx_epochs_done = checkpoint['tfx_epochs_done']\n","    print(f\"tfx_epochs_done is {tfx_epochs_done}\")\n","    model = ResCsNet(N, M)\n","    model.load_state_dict(checkpoint['state_dict'])\n","    model.train()\n","    model.cuda()\n","    \n","    optimizer = optim.Adam(model.parameters())\n","    optimizer.load_state_dict(checkpoint['optimizer'])\n","    \n","    \n","    lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.2, patience=10, threshold=5e-3 ,verbose=True)\n","    lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n","\n","    for state in optimizer.state.values():\n","        for k, v in state.items():\n","            if isinstance(v, torch.Tensor):\n","    #             print(\"copy to cuda\")\n","                state[k] = v.cuda()\n","\n","    print(\"Done\")"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"W-7JjRN0TnEe","colab":{}},"source":["# (re)setting learning rate if needed\n","for g in optimizer.param_groups:\n","    g['lr'] = 1e-5"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"WJ3pL8kkA46a","colab":{}},"source":["print(\"The current lrs:\")\n","for g in optimizer.param_groups:\n","    print(g['lr'])"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"5yYCnafRZx8g"},"source":["## Train and validation functions\n","This section is for trainning. Please ensure that you have `pytorch_dwt_ssim.py` and `trainning_func.py` placed along with this notebook in the same directory before you run the cells in this section.\n","\n","This section is not needed if you just want to test a model."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"wFGR4e4hZx8n","colab":{}},"source":["# The trainning function needs some logging\n","import logging\n","logging.basicConfig(format=\"[%(asctime)s] %(message)s\", filename=exp_name+\"_trainning.log\",level=logging.INFO)"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"HvAHlRtE24zj","colab":{}},"source":["class FrobeniusLoss(nn.Module):\n","    def __init__(self):\n","        super().__init__()\n","\n","    def forward(self, x2, x1, eps=1e-8):\n","        '''\n","        x1, x2: both of shape [N, C, H, W]. x1 is the source\n","        '''\n","\n","        diff = x1 - x2\n","        num = torch.norm(diff, p='fro')\n","        den = torch.norm(x1, p='fro') + eps\n","\n","        frob = num / den\n","\n","        return frob"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"f6GT7Z4pZx8t","scrolled":true,"colab":{}},"source":["from pathlib import Path\n","from pydrive.auth import GoogleAuth\n","from pydrive.drive import GoogleDrive\n","from google.colab import auth\n","from oauth2client.client import GoogleCredentials\n","from trainning_func import get_evaluation\n","from pytorch_dwt_ssim import DWT_SSIM\n","\n","stop_file = Path('_stop')\n","\n","\n","def save_checkpoint(model, optimizer, lr_scheduler, tfx_steps, tfx_epochs_done, ckpt_name):\n","    # Save the state model, just in case\n","    logging.info(\"Saving the state of model\")\n","    state = {\n","        'tfx_steps': tfx_steps,\n","        'tfx_epochs_done': tfx_epochs_done,\n","        'state_dict': model.state_dict(),\n","        'optimizer' : optimizer.state_dict(),\n","        'lr_scheduler': lr_scheduler.state_dict()\n","    }\n","    torch.save(state, ckpt_name)\n","    logging.info(\"State saving done, uploading to google drive\")\n","    \n","    try:\n","        gauth = GoogleAuth()\n","        gauth.LoadCredentialsFile(\"mycreds.txt\")\n","        if gauth.credentials is None:\n","            # Authenticate if they're not there\n","            # self.gauth.LocalWebserverAuth()\n","            print(\"[!] No creds saved\")\n","            auth.authenticate_user()\n","            gauth.credentials = GoogleCredentials.get_application_default()\n","        elif gauth.access_token_expired:\n","            # Refresh them if expired\n","            print(\"[!] Token expired\")\n","            gauth.Refresh()\n","        else:\n","            # Initialize the saved creds\n","#             print(\"Initialize the saved creds\")\n","            gauth.Authorize()\n","        # Save the current credentials to a file\n","        gauth.SaveCredentialsFile(\"mycreds.txt\")\n","        the_drive = GoogleDrive(gauth)\n","        \n","        id_file = Path('_id')\n","        if id_file.exists():\n","            fileid = id_file.read_text()\n","            _gfile_ckpt = the_drive.CreateFile({'id': fileid})\n","            _gfile_ckpt.SetContentFile(ckpt_name)\n","            _gfile_ckpt.Upload()\n","        else:\n","            _gfile_ckpt = the_drive.CreateFile()      \n","            _gfile_ckpt.SetContentFile(ckpt_name)\n","            _gfile_ckpt.Upload()\n","            id_file.write_text(_gfile_ckpt['id']) \n","        \n","    except:\n","        print(\"??? Some error occured when trying to upload to gdrive\")\n","#         raise\n","\n","        \n","\n","def train(model, optimizer, lr_scheduler,\n","          fn_mse=FrobeniusLoss(), fn_cwssim=DWT_SSIM(J=3, wave='haar'),\n","          mse_weight=0.3, cwssim_weight=0.7,\n","          epochs=1, logdir=None, print_every=10, tfx_steps=0, tfx_epochs_done=0, device=torch.device('cuda'),\n","          ckpt_name=\"checkpoint.pt\"):\n","    \"\"\"\n","    Train a model\n","\n","    Inputs:\n","    - model: A PyTorch Module giving the model to train.\n","    - optimizer: An Optimizer object we will use to train the model\n","    - lr_scheduler: learning rate scheduler\n","    - fn_mse, fn_cwssim: loss functions for L2 loss and wavelet loss\n","    - mse_weight, cwssim_weight: weights indicating how important they contribute to the total loss\n","    - epochs: A Python integer giving the number of epochs to train for\n","    - logdir: string. Used to specific the logdir of tensorboard\n","    - print_every: after print_every epochs this function will report to logging\n","    - tfx_steps, tfx_epochs_done: helps tensorboardX summary writer find what current step and epoch is\n","    - device: torch.device('cuda') or torch.device('cpu')\n","\n","    Returns:\n","    - tfx_steps: the end of the tfx_steps\n","    \"\"\"\n","    try:\n","        writer = SummaryWriter(log_dir=logdir)\n","        print(f\"Run `tensorboard --logdir={logdir} --host=127.0.0.1` to visualize in realtime\")\n","\n","        fn_mse = fn_mse.to(device=device)\n","        fn_cwssim = fn_cwssim.to(device=device)\n","\n","        # PSNR scheduler\n","#         lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.2, patience=10, threshold=5e-3 ,verbose=True)\n","        # CW-SSIM scheduler\n","    #     lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', patience=10,\n","    #                                                         verbose=True)\n","        for e in range(epochs):\n","            if stop_file.exists():\n","                print(\"Stop file found. Will stop trainning now\")\n","                save_checkpoint(model,\n","                                optimizer,\n","                                lr_scheduler,\n","                                tfx_steps=tfx_steps,\n","                                tfx_epochs_done=tfx_epochs_done,\n","                                ckpt_name=ckpt_name)\n","                break\n","\n","            model.train()  # ensure the model is in training mode\n","            logging.info('-----------------------------')\n","            logging.info(f'* epoch {tfx_epochs_done}')\n","            for t, original_im in enumerate(loader_train):\n","                original_im = original_im.to(device=device)\n","\n","                recovered_im = model(original_im)\n","                mse_loss = fn_mse(recovered_im, original_im)\n","\n","                cwssim_loss = 1 - fn_cwssim(recovered_im, original_im)\n","\n","                # construct the total loss\n","                loss = mse_weight*mse_loss + cwssim_weight*cwssim_loss\n","\n","                writer.add_scalars('train/loss',\n","                                   {\n","                                       'mse_loss.item()': mse_loss.item(),\n","                                       'cwssim_loss.item()': cwssim_loss.item(),\n","                                       'loss.item()': loss.item()\n","                                   },\n","                                   tfx_steps)\n","\n","                optimizer.zero_grad()\n","                loss.backward()\n","\n","                optimizer.step()\n","\n","                if t % int(print_every) == 0:\n","                    logging.info('Iteration %d/%d, loss = %.4f' % (t, len(loader_train) , loss.item()))\n","\n","                tfx_steps += 1\n","            #end for\n","\n","            # after the end of each epoch\n","            ## increment tfx_epochs_done counter\n","            tfx_epochs_done += 1\n","            ## check the performance of the model\n","            logging.info(\"Checking on subtrain and validation set...\")\n","            subtrain_psnr, val_psnr, subtrain_mix_psnr, val_mix_psnr = get_evaluation(\n","                model,\n","                loader_val,\n","                loader_subtrain,\n","                device,\n","                fn_mse=nn.MSELoss(),\n","                fn_cwssim=fn_cwssim,\n","                mse_weight=mse_weight,\n","                cwssim_weight=cwssim_weight)\n","            logging.info(f\"Average PSNR for subtrain set is {subtrain_psnr} dB\")\n","            logging.info(f\"Average PSNR for validation set is {val_psnr} dB\")\n","\n","            logging.info(f\"Average mixed gain for subtrain set is {subtrain_mix_psnr} dB\")\n","            logging.info(f\"Average mixed gain for validation set is {val_mix_psnr} dB\")\n","\n","    #         loss_subtrain = mse_weight*subtrain_avg_mse + cwssim_weight*subtrain_avg_cwssim_loss\n","    #         loss_val = mse_weight*val_avg_mse + cwssim_weight*val_avg_cwssim_loss\n","            if lr_scheduler is not None:\n","                lr_scheduler.step(val_mix_psnr)  # check loss and determine if the lr should be decreased\n","\n","            writer.add_scalars('train/val_evaluation', \n","                               {\n","                                   'subtrain_psnr': subtrain_psnr.item(),\n","                                   'val_psnr': val_psnr.item(),\n","                                   'subtrain_mix_psnr': subtrain_mix_psnr.item(),\n","                                   'val_mix_psnr': val_mix_psnr.item()\n","                               },\n","                               tfx_epochs_done\n","                              )\n","            # Save the state model, just in case\n","            logging.info(\"Saving the state of model\")\n","            save_checkpoint(model,\n","                            optimizer,\n","                            lr_scheduler,\n","                            tfx_steps=tfx_steps,\n","                            tfx_epochs_done=tfx_epochs_done,\n","                            ckpt_name=ckpt_name)\n","\n","        #end for\n","\n","        writer.close()  # tensorboardX writer\n","        return tfx_steps, tfx_epochs_done\n","    \n","    except (KeyboardInterrupt, SystemExit):\n","        print(\"KeyboardInterrupt: save the state of model\")\n","        save_checkpoint(model,\n","                        optimizer,\n","                        lr_scheduler,\n","                        tfx_steps=tfx_steps,\n","                        tfx_epochs_done=tfx_epochs_done,\n","                        ckpt_name=ckpt_name)\n","        return tfx_steps, tfx_epochs_done\n","    except:\n","        print(\"Emergency: save the state of model\")\n","        save_checkpoint(model,\n","                        optimizer,\n","                        lr_scheduler,\n","                        tfx_steps=tfx_steps,\n","                        tfx_epochs_done=tfx_epochs_done,\n","                        ckpt_name=ckpt_name)\n","        raise"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"DZKtURxBHjbm","colab":{}},"source":["# old style train, mse only\n","def get_avg_psnr(model, loader_val, loader_subtrain, device, print_every=10):\n","    model.eval()  # ensure the model is in evaluation mode\n","    subtrain_avg_psnr = 0\n","    val_avg_psnr = 0\n","\n","    with torch.no_grad():\n","        for t, val_im in enumerate(loader_val):\n","            if t % int(print_every) == 0:\n","                logging.info(f\"checked {t}/{len(loader_val)} in loader_val\")\n","            val_original = val_im.to(device)\n","            val_recovered = model(val_original)\n","            val_mse = F.mse_loss(val_recovered, val_original)\n","            \n","            # PSNR\n","            val_psnr = 10 * np.log10(1 / val_mse.item())\n","            val_avg_psnr += val_psnr\n","            \n","        val_avg_psnr /= len(loader_val)\n","\n","\n","        for t, subtrain_im in enumerate(loader_subtrain):\n","            if t % int(print_every) == 0:\n","                logging.info(f\"checked {t}/{len(loader_subtrain)} in loader_subtrain\")\n","            subtrain_original = subtrain_im.to(device)\n","            subtrain_recovered = model(subtrain_original)\n","            subtrain_mse = F.mse_loss(subtrain_recovered, subtrain_original)\n","            \n","            # PSNR\n","            subtrain_psnr = 10 * np.log10(1 / subtrain_mse.item())\n","            subtrain_avg_psnr += subtrain_psnr\n","            \n","        subtrain_avg_psnr /= len(loader_subtrain)\n","    \n","    return subtrain_avg_psnr, val_avg_psnr\n","\n","#############################################################\n","\n","def train_oldstyle(model, optimizer, lr_scheduler,\n","          epochs=1, logdir=None, print_every=10, tfx_steps=0, tfx_epochs_done=0, device=torch.device('cuda'),\n","          ckpt_name=\"checkpoint.pt\"):\n","    \"\"\"\n","    Train a model\n","\n","    Inputs:\n","    - model: A PyTorch Module giving the model to train.\n","    - optimizer: An Optimizer object we will use to train the model\n","    - lr_scheduler: learning rate scheduler\n","    - epochs: A Python integer giving the number of epochs to train for\n","    - logdir: string. Used to specific the logdir of tensorboard\n","\n","    Returns:\n","    - tfx_steps: the end of the tfx_steps\n","    \"\"\"\n","    try:\n","        writer = SummaryWriter(log_dir=logdir)\n","        print(f\"Run `tensorboard --logdir={logdir} --host=127.0.0.1` to visualize in realtime\")\n","\n","        fn_mse = nn.MSELoss()\n","        fn_mse = fn_mse.to(device=device)\n","        \n","        for e in range(epochs):\n","            if stop_file.exists():\n","                print(\"Stop file found. Will stop trainning now\")\n","                save_checkpoint(model,\n","                                optimizer,\n","                                lr_scheduler,\n","                                tfx_steps=tfx_steps,\n","                                tfx_epochs_done=tfx_epochs_done,\n","                                ckpt_name=ckpt_name)\n","                break\n","\n","            model.train()  # ensure the model is in training mode\n","            logging.info('-----------------------------')\n","            logging.info(f'* epoch {tfx_epochs_done}')\n","            for t, original_im in enumerate(loader_train):\n","                original_im = original_im.to(device=device)\n","\n","                recovered_im = model(original_im)\n","                mse_loss = fn_mse(recovered_im, original_im)\n","\n","                # construct the total loss\n","                loss = mse_loss\n","\n","                writer.add_scalars('train/loss',\n","                                   {\n","                                       'mse_loss.item()': mse_loss.item(),\n","                                       'loss.item()': loss.item()\n","                                   },\n","                                   tfx_steps)\n","\n","                optimizer.zero_grad()\n","                loss.backward()\n","\n","                optimizer.step()\n","\n","                if t % int(print_every) == 0:\n","                    logging.info('Iteration %d/%d, loss = %.4f' % (t, len(loader_train) , loss.item()))\n","\n","                tfx_steps += 1\n","            #end for\n","\n","            # after the end of each epoch\n","            ## increment tfx_epochs_done counter\n","            tfx_epochs_done += 1\n","            ## check the performance of the model\n","            logging.info(\"Checking on subtrain and validation set...\")\n","            subtrain_psnr, val_psnr = get_avg_psnr(model, loader_val, loader_subtrain, device)\n","            \n","            logging.info(f\"Average PSNR for subtrain set is {subtrain_psnr} dB\")\n","            logging.info(f\"Average PSNR for validation set is {val_psnr} dB\")\n","\n","            if lr_scheduler is not None:\n","                lr_scheduler.step(val_psnr)  # check loss and determine if the lr should be decreased\n","\n","            writer.add_scalars('train/val_evaluation', \n","                               {\n","                                   'subtrain_psnr': subtrain_psnr.item(),\n","                                   'val_psnr': val_psnr.item(),\n","                               },\n","                               tfx_epochs_done\n","                              )\n","            # Save the state model, just in case\n","            logging.info(\"Saving the state of model\")\n","            save_checkpoint(model,\n","                            optimizer,\n","                            lr_scheduler,\n","                            tfx_steps=tfx_steps,\n","                            tfx_epochs_done=tfx_epochs_done,\n","                            ckpt_name=ckpt_name)\n","\n","        #end for\n","\n","        writer.close()  # tensorboardX writer\n","        return tfx_steps, tfx_epochs_done\n","    \n","    except (KeyboardInterrupt, SystemExit):\n","        print(\"KeyboardInterrupt: save the state of model\")\n","        save_checkpoint(model,\n","                        optimizer,\n","                        lr_scheduler,\n","                        tfx_steps=tfx_steps,\n","                        tfx_epochs_done=tfx_epochs_done,\n","                        ckpt_name=ckpt_name)\n","        return tfx_steps, tfx_epochs_done\n","    except:\n","        print(\"Emergency: save the state of model\")\n","        save_checkpoint(model,\n","                        optimizer,\n","                        lr_scheduler,\n","                        tfx_steps=tfx_steps,\n","                        tfx_epochs_done=tfx_epochs_done,\n","                        ckpt_name=ckpt_name)\n","        raise"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"E0Sy07uiZx83"},"source":["## Time consuming part ...\n","It is a good idea to train with `train_oldstyle` which only uses MSE loss as the criterion, then switch to `train` which combines l2 loss and DW-SSIM loss to finetune the model. Running wavelet code is still slower even on GPU. By using this trainning scheme you can save some time."]},{"cell_type":"code","metadata":{"colab_type":"code","id":"RjmLFXPGZx84","scrolled":true,"colab":{}},"source":["learning_rate = 5e-4\n","tfx_steps = 0\n","tfx_epochs_done = 0\n","model = model.to(device=device)  # move to proper device before constructing the optimizer\n","optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n","lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.2, patience=10, threshold=5e-3 ,verbose=True)\n","\n","# train_oldstyle only incorporates L2 loss (nn.MSELoss())\n","# to use wavelet loss function, use the train function to train\n","tfx_steps, tfx_epochs_done = train_oldstyle(model, optimizer, lr_scheduler,\n","                                            epochs=49, logdir='runs/' + exp_name,\n","                                            tfx_steps=tfx_steps, tfx_epochs_done=tfx_epochs_done, device=device,\n","                                            ckpt_name=exp_name+\"_checkpoint.pth\")"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"Jw62TncGlrsF","colab":{}},"source":["tfx_steps, tfx_epochs_done = train_oldstyle(model, optimizer, lr_scheduler,\n","                                            epochs=35, logdir='runs/' + exp_name,\n","                                            tfx_steps=tfx_steps, tfx_epochs_done=tfx_epochs_done, device=device,\n","                                            ckpt_name=exp_name+\"_checkpoint.pth\")"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"ZfTdO94Yny66","colab":{}},"source":["tfx_steps, tfx_epochs_done = train(model, optimizer, lr_scheduler,\n","                                   mse_weight=0.5, cwssim_weight=0.5,\n","                                   epochs=22, logdir='runs/' + exp_name,\n","                                   tfx_steps=tfx_steps, tfx_epochs_done=tfx_epochs_done, print_every=10, device=device,\n","                                   ckpt_name=exp_name+\"_checkpoint.pth\")"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"lCzX22RXnFAq","colab":{}},"source":["tfx_steps, tfx_epochs_done = train(model, optimizer, lr_scheduler,\n","                                   mse_weight=0.5, cwssim_weight=0.5,\n","                                   epochs=22, logdir='runs/' + exp_name,\n","                                   tfx_steps=tfx_steps, tfx_epochs_done=tfx_epochs_done, print_every=10, device=device,\n","                                   ckpt_name=exp_name+\"_checkpoint.pth\")"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"uslYGplXZx9G","scrolled":true,"colab":{}},"source":["# print(\"The current lrs:\")\n","for g in optimizer.param_groups:\n","    g['lr'] = 1e-6"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"XU4blAJ6TtwW","colab":{}},"source":["tfx_steps, tfx_epochs_done = train(model, optimizer, lr_scheduler,\n","                                   mse_weight=0.8, cwssim_weight=0.2,\n","                                   epochs=22, logdir='runs/' + exp_name,\n","                                   tfx_steps=tfx_steps, tfx_epochs_done=tfx_epochs_done, print_every=10, device=device,\n","                                   ckpt_name=exp_name+\"_checkpoint.pth\")"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"0LAJRc1fbuUb","colab":{}},"source":["lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.2, patience=10, threshold=5e-3 ,verbose=True)\n","tfx_steps, tfx_epochs_done = train(model, optimizer, lr_scheduler,\n","                                   mse_weight=0.8, cwssim_weight=0.2,\n","                                   epochs=13, logdir='runs/' + exp_name,\n","                                   tfx_steps=tfx_steps, tfx_epochs_done=tfx_epochs_done, print_every=10, device=device,\n","                                   ckpt_name=exp_name+\"_checkpoint.pth\")"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"eNzloahnsl7M","colab":{}},"source":["tfx_steps, tfx_epochs_done = train(model, optimizer, lr_scheduler,\n","                                   mse_weight=0.5, cwssim_weight=0.5,\n","                                   epochs=13, logdir='runs/' + exp_name,\n","                                   tfx_steps=tfx_steps, tfx_epochs_done=tfx_epochs_done, print_every=10, device=device,\n","                                   ckpt_name=exp_name+\"_checkpoint.pth\")"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"smcgBFP5Zx9J"},"source":["# Save the model again (just in case)"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"osq5b2MrZx9K","colab":{}},"source":["print(\"Saving the state of model\")\n","state = {\n","    'tfx_steps': tfx_steps,\n","    'tfx_epochs_done': tfx_epochs_done,\n","    'state_dict': model.state_dict(),\n","    'optimizer' : optimizer.state_dict(),\n","}\n","torch.save(state, exp_name+\"_checkpoint.bak.pth\")\n","\n","print(\"Uploading to google drive\")\n","gfile_ckpt.SetContentFile(ckpt_name)\n","gfile_ckpt.Upload()\n","\n","print(\"Done\")"],"execution_count":0,"outputs":[]},{"cell_type":"markdown","metadata":{"colab_type":"text","id":"R05lsaxfZx9M"},"source":["# Testing"]},{"cell_type":"code","metadata":{"colab_type":"code","id":"aB2Lq0SaZx9N","scrolled":true,"colab":{}},"source":["# Visualize the recovered and original image\n","def denormalize(im):\n","    if im.shape[2] == 1: # grayscale\n","        mean, std = (0.5,), (0.5,)\n","        im = im.reshape(im.shape[0], im.shape[1])\n","        im = im * std[0] + mean[0]\n","    else: # rgb\n","        mean, std = (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)\n","        for ch in range(3):\n","            im[:,:,ch] = im[:,:,ch] * std[ch] + mean[ch]\n","    return im\n","\n","# visualize recovered image\n","def vis_show(model, idx=0):\n","    model = model.to(device=device)\n","    model.eval()\n","\n","    test_iter = iter(loader_test)\n","    _y = test_iter.next()\n","    print(_y.shape)\n","\n","    with torch.no_grad():\n","        _y = _y.to(device=torch.device(\"cuda\"))\n","        restored_im = model(_y)\n","    print(restored_im.shape)\n","    _im = restored_im[idx].cpu().numpy()\n","    _im = _im.transpose(1, 2, 0).astype(np.float)\n","    _ori = _y[idx].cpu().numpy()\n","    _ori = _ori.transpose(1, 2, 0).astype(np.float)\n","    # mean, std = (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)\n","\n","    # _im[_im < -1] = -1.0\n","    # _im[_im > 1] = 1.0\n","\n","    mse = np.mean((_im - _ori) ** 2 )\n","    psnr = 10 * np.log10(1/mse)\n","    print(f\"mse is {mse}, psnr is {psnr} dB\")\n","\n","    _im = denormalize(_im)\n","    _ori = denormalize(_ori)\n","    # im_mx = np.amax(_im)\n","    # im_mn = np.amin(_im)\n","    # _im = _im / (im_mx-im_mn) * 1.0\n","\n","    # print(f'_im is {_im}')\n","    # print(f'_ori is {_ori}')\n","\n","\n","    print(f\"the original image: mx = {np.amax(_ori)}, mn={np.amin(_ori)}\")\n","    print(f\"the restored image: mx = {np.amax(_im)}, mn={np.amin(_im)}\")\n","    plt.subplot(1,2,1)\n","    plt.imshow(_ori, cmap=\"gray\"); plt.title('original')\n","    plt.grid(False)\n","    plt.subplot(1,2,2)\n","    plt.imshow(_im, cmap=\"gray\"); plt.title('restored')\n","    plt.grid(False)\n","    plt.show()\n","    \n","vis_show(model, 6)"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"He5SnneIZx9R","scrolled":true,"colab":{}},"source":["# Calculate PSNR\n","device = torch.device('cuda')\n","model.to(device)\n","avg_psnr = 0\n","_psnrs = []\n","with torch.no_grad():\n","    model.eval()\n","    for t,batch in enumerate(loader_test):\n","        if t % 10 == 0: print(f\"{t}/{len(loader_test)}\")\n","        \n","        original = batch.to(device)\n","        recovered = model(original)\n","            \n","        recovered = denormalize(recovered)\n","        original = denormalize(original)\n","        diff = recovered - original\n","        rmse = np.sqrt( np.mean(diff ** 2 ) )\n","        psnr = 20 * np.log10(1/rmse)\n","#         mse = F.mse_loss(recovered, original)\n","#         psnr = 10 * np.log10(1 / mse.item())\n","        _psnrs.append(psnr)\n","        avg_psnr += psnr\n","print(\"===> Avg. PSNR: {:.4f} dB\".format(avg_psnr / len(loader_test)))\n","\n","psnrs = np.array(_psnrs)\n","plt.hist(psnrs, bins=15); plt.xlabel('PSNR/dB'); plt.ylabel('Number of samples')\n","plt.show()"],"execution_count":0,"outputs":[]},{"cell_type":"code","metadata":{"colab_type":"code","id":"IgwBP2YPAu3F","colab":{}},"source":[""],"execution_count":0,"outputs":[]}]}