first commit

.gitignore

+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# PyTorch
+*.pt
+*.pdf
+*.png
+*.txt
+*.swp
+.vscode
+
+# Datasets
+data/
+test/
\ No newline at end of file

LICENSE

+MIT License
+
+Copyright (c) 2018 Sanghyun Son
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

+**About PyTorch 1.2.0**
+  * Now the master branch supports PyTorch 1.2.0 by default.
+  * Due to the serious version problem (especially torch.utils.data.dataloader), MDSR functions are temporarily disabled. If you have to train/evaluate the MDSR model, please use legacy branches.
+
+# EDSR-PyTorch
+
+**About PyTorch 1.1.0**
+  * There have been minor changes with the 1.1.0 update. Now we support PyTorch 1.1.0 by default, and please use the legacy branch if you prefer older version.
+
+![](/figs/main.png)
+
+This repository is an official PyTorch implementation of the paper **"Enhanced Deep Residual Networks for Single Image Super-Resolution"** from **CVPRW 2017, 2nd NTIRE**.
+You can find the original code and more information from [here](https://github.com/LimBee/NTIRE2017).
+
+If you find our work useful in your research or publication, please cite our work:
+
+[1] Bee Lim, Sanghyun Son, Heewon Kim, Seungjun Nah, and Kyoung Mu Lee, **"Enhanced Deep Residual Networks for Single Image Super-Resolution,"** <i>2nd NTIRE: New Trends in Image Restoration and Enhancement workshop and challenge on image super-resolution in conjunction with **CVPR 2017**. </i> [[PDF](http://openaccess.thecvf.com/content_cvpr_2017_workshops/w12/papers/Lim_Enhanced_Deep_Residual_CVPR_2017_paper.pdf)] [[arXiv](https://arxiv.org/abs/1707.02921)] [[Slide](https://cv.snu.ac.kr/research/EDSR/Presentation_v3(release).pptx)]
+```
+@InProceedings{Lim_2017_CVPR_Workshops,
+  author = {Lim, Bee and Son, Sanghyun and Kim, Heewon and Nah, Seungjun and Lee, Kyoung Mu},
+  title = {Enhanced Deep Residual Networks for Single Image Super-Resolution},
+  booktitle = {The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) Workshops},
+  month = {July},
+  year = {2017}
+}
+```
+We provide scripts for reproducing all the results from our paper. You can train your model from scratch, or use a pre-trained model to enlarge your images.
+
+**Differences between Torch version**
+* Codes are much more compact. (Removed all unnecessary parts.)
+* Models are smaller. (About half.)
+* Slightly better performances.
+* Training and evaluation requires less memory.
+* Python-based.
+
+## Dependencies
+* Python 3.6
+* PyTorch >= 1.0.0
+* numpy
+* skimage
+* **imageio**
+* matplotlib
+* tqdm
+* cv2 >= 3.xx (Only if you want to use video input/output)
+
+## Code
+Clone this repository into any place you want.
+```bash
+git clone https://github.com/thstkdgus35/EDSR-PyTorch
+cd EDSR-PyTorch
+```
+
+## Quickstart (Demo)
+You can test our super-resolution algorithm with your images. Place your images in ``test`` folder. (like ``test/<your_image>``) We support **png** and **jpeg** files.
+
+Run the script in ``src`` folder. Before you run the demo, please uncomment the appropriate line in ```demo.sh``` that you want to execute.
+```bash
+cd src       # You are now in */EDSR-PyTorch/src
+sh demo.sh
+```
+
+You can find the result images from ```experiment/test/results``` folder.
+
+| Model | Scale | File name (.pt) | Parameters | ****PSNR** |
+|  ---  |  ---  | ---       | ---        | ---  |
+| **EDSR** | 2 | EDSR_baseline_x2 | 1.37 M | 34.61 dB |
+| | | *EDSR_x2 | 40.7 M | 35.03 dB |
+| | 3 | EDSR_baseline_x3 | 1.55 M | 30.92 dB |
+| | | *EDSR_x3 | 43.7 M | 31.26 dB |
+| | 4 | EDSR_baseline_x4 | 1.52 M | 28.95 dB |
+| | | *EDSR_x4 | 43.1 M | 29.25 dB |
+| **MDSR** | 2 | MDSR_baseline | 3.23 M | 34.63 dB |
+| | | *MDSR | 7.95 M| 34.92 dB |
+| | 3 | MDSR_baseline | | 30.94 dB |
+| | | *MDSR | | 31.22 dB |
+| | 4 | MDSR_baseline | | 28.97 dB |
+| | | *MDSR | | 29.24 dB |
+
+*Baseline models are in ``experiment/model``. Please download our final models from [here](https://cv.snu.ac.kr/research/EDSR/model_pytorch.tar) (542MB)
+**We measured PSNR using DIV2K 0801 ~ 0900, RGB channels, without self-ensemble. (scale + 2) pixels from the image boundary are ignored.
+
+You can evaluate your models with widely-used benchmark datasets:
+
+[Set5 - Bevilacqua et al. BMVC 2012](http://people.rennes.inria.fr/Aline.Roumy/results/SR_BMVC12.html),
+
+[Set14 - Zeyde et al. LNCS 2010](https://sites.google.com/site/romanzeyde/research-interests),
+
+[B100 - Martin et al. ICCV 2001](https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/bsds/),
+
+[Urban100 - Huang et al. CVPR 2015](https://sites.google.com/site/jbhuang0604/publications/struct_sr).
+
+For these datasets, we first convert the result images to YCbCr color space and evaluate PSNR on the Y channel only. You can download [benchmark datasets](https://cv.snu.ac.kr/research/EDSR/benchmark.tar) (250MB). Set ``--dir_data <where_benchmark_folder_located>`` to evaluate the EDSR and MDSR with the benchmarks.
+
+You can download some results from [here](https://cv.snu.ac.kr/research/EDSR/result_image/edsr-results.tar).
+The link contains **EDSR+_baseline_x4** and **EDSR+_x4**.
+Otherwise, you can easily generate result images with ``demo.sh`` scripts.
+
+## How to train EDSR and MDSR
+We used [DIV2K](http://www.vision.ee.ethz.ch/%7Etimofter/publications/Agustsson-CVPRW-2017.pdf) dataset to train our model. Please download it from [here](https://cv.snu.ac.kr/research/EDSR/DIV2K.tar) (7.1GB).
+
+Unpack the tar file to any place you want. Then, change the ```dir_data``` argument in ```src/option.py``` to the place where DIV2K images are located.
+
+We recommend you to pre-process the images before training. This step will decode all **png** files and save them as binaries. Use ``--ext sep_reset`` argument on your first run. You can skip the decoding part and use saved binaries with ``--ext sep`` argument.
+
+If you have enough RAM (>= 32GB), you can use ``--ext bin`` argument to pack all DIV2K images in one binary file.
+
+You can train EDSR and MDSR by yourself. All scripts are provided in the ``src/demo.sh``. Note that EDSR (x3, x4) requires pre-trained EDSR (x2). You can ignore this constraint by removing ```--pre_train <x2 model>``` argument.
+
+```bash
+cd src       # You are now in */EDSR-PyTorch/src
+sh demo.sh
+```
+
+**Update log**
+* Jan 04, 2018
+  * Many parts are re-written. You cannot use previous scripts and models directly.
+  * Pre-trained MDSR is temporarily disabled.
+  * Training details are included.
+
+* Jan 09, 2018
+  * Missing files are included (```src/data/MyImage.py```).
+  * Some links are fixed.
+
+* Jan 16, 2018
+  * Memory efficient forward function is implemented.
+  * Add --chop_forward argument to your script to enable it.
+  * Basically, this function first split a large image to small patches. Those images are merged after super-resolution. I checked this function with 12GB memory, 4000 x 2000 input image in scale 4. (Therefore, the output will be 16000 x 8000.)
+
+* Feb 21, 2018
+  * Fixed the problem when loading pre-trained multi-GPU model.
+  * Added pre-trained scale 2 baseline model.
+  * This code now only saves the best-performing model by default. For MDSR, 'the best' can be ambiguous. Use --save_models argument to keep all the intermediate models.
+  * PyTorch 0.3.1 changed their implementation of DataLoader function. Therefore, I also changed my implementation of MSDataLoader. You can find it on feature/dataloader branch.
+
+* Feb 23, 2018
+  * Now PyTorch 0.3.1 is a default. Use legacy/0.3.0 branch if you use the old version.
+  * With a new ``src/data/DIV2K.py`` code, one can easily create new data class for super-resolution.
+  * New binary data pack. (Please remove the ``DIV2K_decoded`` folder from your dataset if you have.)
+  * With ``--ext bin``, this code will automatically generate and saves the binary data pack that corresponds to previous ``DIV2K_decoded``. (This requires huge RAM (~45GB, Swap can be used.), so please be careful.)
+  * If you cannot make the binary pack, use the default setting (``--ext img``).
+
+  * Fixed a bug that PSNR in the log and PSNR calculated from the saved images does not match.
+  * Now saved images have better quality! (PSNR is ~0.1dB higher than the original code.)
+  * Added performance comparison between Torch7 model and PyTorch models.
+
+* Mar 5, 2018
+  * All baseline models are uploaded.
+  * Now supports half-precision at test time. Use ``--precision half``  to enable it. This does not degrade the output images.
+
+* Mar 11, 2018
+  * Fixed some typos in the code and script.
+  * Now --ext img is default setting. Although we recommend you to use --ext bin when training, please use --ext img when you use --test_only.
+  * Skip_batch operation is implemented. Use --skip_threshold argument to skip the batch that you want to ignore. Although this function is not exactly the same with that of Torch7 version, it will work as you expected.
+
+* Mar 20, 2018
+  * Use ``--ext sep-reset`` to pre-decode large png files. Those decoded files will be saved to the same directory with DIV2K png files. After the first run, you can use ``--ext sep`` to save time.
+  * Now supports various benchmark datasets. For example, try ``--data_test Set5`` to test your model on the Set5 images.
+  * Changed the behavior of skip_batch.
+
+* Mar 29, 2018
+  * We now provide all models from our paper.
+  * We also provide ``MDSR_baseline_jpeg`` model that suppresses JPEG artifacts in the original low-resolution image. Please use it if you have any trouble.
+  * ``MyImage`` dataset is changed to ``Demo`` dataset. Also, it works more efficient than before.
+  * Some codes and script are re-written.
+
+* Apr 9, 2018
+  * VGG and Adversarial loss is implemented based on [SRGAN](http://openaccess.thecvf.com/content_cvpr_2017/papers/Ledig_Photo-Realistic_Single_Image_CVPR_2017_paper.pdf). [WGAN](https://arxiv.org/abs/1701.07875) and [gradient penalty](https://arxiv.org/abs/1704.00028) are also implemented, but they are not tested yet.
+  * Many codes are refactored. If there exists a bug, please report it.
+  * [D-DBPN](https://arxiv.org/abs/1803.02735) is implemented. The default setting is D-DBPN-L.
+
+* Apr 26, 2018
+  * Compatible with PyTorch 0.4.0
+  * Please use the legacy/0.3.1 branch if you are using the old version of PyTorch.
+  * Minor bug fixes
+
+* July 22, 2018
+  * Thanks for recent commits that contains RDN and RCAN. Please see ``code/demo.sh`` to train/test those models.
+  * Now the dataloader is much stable than the previous version. Please erase ``DIV2K/bin`` folder that is created before this commit. Also, please avoid using ``--ext bin`` argument. Our code will automatically pre-decode png images before training. If you do not have enough spaces(~10GB) in your disk, we recommend ``--ext img``(But SLOW!).
+
+* Oct 18, 2018
+  * with ``--pre_train download``, pretrained models will be automatically downloaded from the server.
+  * Supports video input/output (inference only). Try with ``--data_test video --dir_demo [video file directory]``.
+
+* About PyTorch 1.0.0
+  * We support PyTorch 1.0.0. If you prefer the previous versions of PyTorch, use legacy branches.
+  * ``--ext bin`` is not supported. Also, please erase your bin files with ``--ext sep-reset``. Once you successfully build those bin files, you can remove ``-reset`` from the argument.

experiment/.gitignore

+*
+!.gitignore
+!/model/*.pt

src/dataloader.py

+import threading
+import random
+
+import torch
+import torch.multiprocessing as multiprocessing
+from torch.utils.data import DataLoader
+from torch.utils.data import SequentialSampler
+from torch.utils.data import RandomSampler
+from torch.utils.data import BatchSampler
+from torch.utils.data import _utils
+from torch.utils.data.dataloader import _DataLoaderIter
+
+from torch.utils.data._utils import collate
+from torch.utils.data._utils import signal_handling
+from torch.utils.data._utils import MP_STATUS_CHECK_INTERVAL
+from torch.utils.data._utils import ExceptionWrapper
+from torch.utils.data._utils import IS_WINDOWS
+from torch.utils.data._utils.worker import ManagerWatchdog
+
+from torch._six import queue
+
+def _ms_loop(dataset, index_queue, data_queue, done_event, collate_fn, scale, seed, init_fn, worker_id):
+    try:
+        collate._use_shared_memory = True
+        signal_handling._set_worker_signal_handlers()
+
+        torch.set_num_threads(1)
+        random.seed(seed)
+        torch.manual_seed(seed)
+
+        data_queue.cancel_join_thread()
+
+        if init_fn is not None:
+            init_fn(worker_id)
+
+        watchdog = ManagerWatchdog()
+
+        while watchdog.is_alive():
+            try:
+                r = index_queue.get(timeout=MP_STATUS_CHECK_INTERVAL)
+            except queue.Empty:
+                continue
+
+            if r is None:
+                assert done_event.is_set()
+                return
+            elif done_event.is_set():
+                continue
+
+            idx, batch_indices = r
+            try:
+                idx_scale = 0
+                if len(scale) > 1 and dataset.train:
+                    idx_scale = random.randrange(0, len(scale))
+                    dataset.set_scale(idx_scale)
+
+                samples = collate_fn([dataset[i] for i in batch_indices])
+                samples.append(idx_scale)
+            except Exception:
+                data_queue.put((idx, ExceptionWrapper(sys.exc_info())))
+            else:
+                data_queue.put((idx, samples))
+                del samples
+
+    except KeyboardInterrupt:
+        pass
+
+class _MSDataLoaderIter(_DataLoaderIter):
+
+    def __init__(self, loader):
+        self.dataset = loader.dataset
+        self.scale = loader.scale
+        self.collate_fn = loader.collate_fn
+        self.batch_sampler = loader.batch_sampler
+        self.num_workers = loader.num_workers
+        self.pin_memory = loader.pin_memory and torch.cuda.is_available()
+        self.timeout = loader.timeout
+
+        self.sample_iter = iter(self.batch_sampler)
+
+        base_seed = torch.LongTensor(1).random_().item()
+
+        if self.num_workers > 0:
+            self.worker_init_fn = loader.worker_init_fn
+            self.worker_queue_idx = 0
+            self.worker_result_queue = multiprocessing.Queue()
+            self.batches_outstanding = 0
+            self.worker_pids_set = False
+            self.shutdown = False
+            self.send_idx = 0
+            self.rcvd_idx = 0
+            self.reorder_dict = {}
+            self.done_event = multiprocessing.Event()
+
+            base_seed = torch.LongTensor(1).random_()[0]
+
+            self.index_queues = []
+            self.workers = []
+            for i in range(self.num_workers):
+                index_queue = multiprocessing.Queue()
+                index_queue.cancel_join_thread()
+                w = multiprocessing.Process(
+                    target=_ms_loop,
+                    args=(
+                        self.dataset,
+                        index_queue,
+                        self.worker_result_queue,
+                        self.done_event,
+                        self.collate_fn,
+                        self.scale,
+                        base_seed + i,
+                        self.worker_init_fn,
+                        i
+                    )
+                )
+                w.daemon = True
+                w.start()
+                self.index_queues.append(index_queue)
+                self.workers.append(w)
+
+            if self.pin_memory:
+                self.data_queue = queue.Queue()
+                pin_memory_thread = threading.Thread(
+                    target=_utils.pin_memory._pin_memory_loop,
+                    args=(
+                        self.worker_result_queue,
+                        self.data_queue,
+                        torch.cuda.current_device(),
+                        self.done_event
+                    )
+                )
+                pin_memory_thread.daemon = True
+                pin_memory_thread.start()
+                self.pin_memory_thread = pin_memory_thread
+            else:
+                self.data_queue = self.worker_result_queue
+
+            _utils.signal_handling._set_worker_pids(
+                id(self), tuple(w.pid for w in self.workers)
+            )
+            _utils.signal_handling._set_SIGCHLD_handler()
+            self.worker_pids_set = True
+
+            for _ in range(2 * self.num_workers):
+                self._put_indices()
+
+
+class MSDataLoader(DataLoader):
+
+    def __init__(self, cfg, *args, **kwargs):
+        super(MSDataLoader, self).__init__(
+            *args, **kwargs, num_workers=cfg.n_threads
+        )
+        self.scale = cfg.scale
+
+    def __iter__(self):
+        return _MSDataLoaderIter(self)
+

src/demo.sh

+# EDSR baseline model (x2) + JPEG augmentation
+python main.py --model EDSR --scale 2 --patch_size 96 --save edsr_baseline_x2 --reset
+#python main.py --model EDSR --scale 2 --patch_size 96 --save edsr_baseline_x2 --reset --data_train DIV2K+DIV2K-Q75 --data_test DIV2K+DIV2K-Q75
+
+# EDSR baseline model (x3) - from EDSR baseline model (x2)
+#python main.py --model EDSR --scale 3 --patch_size 144 --save edsr_baseline_x3 --reset --pre_train [pre-trained EDSR_baseline_x2 model dir]
+
+# EDSR baseline model (x4) - from EDSR baseline model (x2)
+#python main.py --model EDSR --scale 4 --save edsr_baseline_x4 --reset --pre_train [pre-trained EDSR_baseline_x2 model dir]
+
+# EDSR in the paper (x2)
+#python main.py --model EDSR --scale 2 --save edsr_x2 --n_resblocks 32 --n_feats 256 --res_scale 0.1 --reset
+
+# EDSR in the paper (x3) - from EDSR (x2)
+#python main.py --model EDSR --scale 3 --save edsr_x3 --n_resblocks 32 --n_feats 256 --res_scale 0.1 --reset --pre_train [pre-trained EDSR model dir]
+
+# EDSR in the paper (x4) - from EDSR (x2)
+#python main.py --model EDSR --scale 4 --save edsr_x4 --n_resblocks 32 --n_feats 256 --res_scale 0.1 --reset --pre_train [pre-trained EDSR_x2 model dir]
+
+# MDSR baseline model
+#python main.py --template MDSR --model MDSR --scale 2+3+4 --save MDSR_baseline --reset --save_models
+
+# MDSR in the paper
+#python main.py --template MDSR --model MDSR --scale 2+3+4 --n_resblocks 80 --save MDSR --reset --save_models
+
+# Standard benchmarks (Ex. EDSR_baseline_x4)
+#python main.py --data_test Set5+Set14+B100+Urban100+DIV2K --data_range 801-900 --scale 4 --pre_train download --test_only --self_ensemble
+
+#python main.py --data_test Set5+Set14+B100+Urban100+DIV2K --data_range 801-900 --scale 4 --n_resblocks 32 --n_feats 256 --res_scale 0.1 --pre_train download --test_only --self_ensemble
+
+# Test your own images
+#python main.py --data_test Demo --scale 4 --pre_train download --test_only --save_results
+
+# Advanced - Test with JPEG images 
+#python main.py --model MDSR --data_test Demo --scale 2+3+4 --pre_train download --test_only --save_results
+
+# Advanced - Training with adversarial loss
+#python main.py --template GAN --scale 4 --save edsr_gan --reset --patch_size 96 --loss 5*VGG54+0.15*GAN --pre_train download
+
+# RDN BI model (x2)
+#python3.6 main.py --scale 2 --save RDN_D16C8G64_BIx2 --model RDN --epochs 200 --batch_size 16 --data_range 801-805 --patch_size 64 --reset
+# RDN BI model (x3)
+#python3.6 main.py --scale 3 --save RDN_D16C8G64_BIx3 --model RDN --epochs 200 --batch_size 16 --data_range 801-805 --patch_size 96 --reset
+# RDN BI model (x4)
+#python3.6 main.py --scale 4 --save RDN_D16C8G64_BIx4 --model RDN --epochs 200 --batch_size 16 --data_range 801-805 --patch_size 128 --reset
+
+# RCAN_BIX2_G10R20P48, input=48x48, output=96x96
+# pretrained model can be downloaded from https://www.dropbox.com/s/mjbcqkd4nwhr6nu/models_ECCV2018RCAN.zip?dl=0
+#python main.py --template RCAN --save RCAN_BIX2_G10R20P48 --scale 2 --reset --save_results --patch_size 96
+# RCAN_BIX3_G10R20P48, input=48x48, output=144x144
+#python main.py --template RCAN --save RCAN_BIX3_G10R20P48 --scale 3 --reset --save_results --patch_size 144 --pre_train ../experiment/model/RCAN_BIX2.pt
+# RCAN_BIX4_G10R20P48, input=48x48, output=192x192
+#python main.py --template RCAN --save RCAN_BIX4_G10R20P48 --scale 4 --reset --save_results --patch_size 192 --pre_train ../experiment/model/RCAN_BIX2.pt
+# RCAN_BIX8_G10R20P48, input=48x48, output=384x384
+#python main.py --template RCAN --save RCAN_BIX8_G10R20P48 --scale 8 --reset --save_results --patch_size 384 --pre_train ../experiment/model/RCAN_BIX2.pt
+

src/loss/__init__.py

+import os
+from importlib import import_module
+
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Loss(nn.modules.loss._Loss):
+    def __init__(self, args, ckp):
+        super(Loss, self).__init__()
+        print('Preparing loss function:')
+
+        self.n_GPUs = args.n_GPUs
+        self.loss = []
+        self.loss_module = nn.ModuleList()
+        for loss in args.loss.split('+'):
+            weight, loss_type = loss.split('*')
+            if loss_type == 'MSE':
+                loss_function = nn.MSELoss()
+            elif loss_type == 'L1':
+                loss_function = nn.L1Loss()
+            elif loss_type.find('VGG') >= 0:
+                module = import_module('loss.vgg')
+                loss_function = getattr(module, 'VGG')(
+                    loss_type[3:],
+                    rgb_range=args.rgb_range
+                )
+            elif loss_type.find('GAN') >= 0:
+                module = import_module('loss.adversarial')
+                loss_function = getattr(module, 'Adversarial')(
+                    args,
+                    loss_type
+                )
+
+            self.loss.append({
+                'type': loss_type,
+                'weight': float(weight),
+                'function': loss_function}
+            )
+            if loss_type.find('GAN') >= 0:
+                self.loss.append({'type': 'DIS', 'weight': 1, 'function': None})
+
+        if len(self.loss) > 1:
+            self.loss.append({'type': 'Total', 'weight': 0, 'function': None})
+
+        for l in self.loss:
+            if l['function'] is not None:
+                print('{:.3f} * {}'.format(l['weight'], l['type']))
+                self.loss_module.append(l['function'])
+
+        self.log = torch.Tensor()
+
+        device = torch.device('cpu' if args.cpu else 'cuda')
+        self.loss_module.to(device)
+        if args.precision == 'half': self.loss_module.half()
+        if not args.cpu and args.n_GPUs > 1:
+            self.loss_module = nn.DataParallel(
+                self.loss_module, range(args.n_GPUs)
+            )
+
+        if args.load != '': self.load(ckp.dir, cpu=args.cpu)
+
+    def forward(self, sr, hr):
+        losses = []
+        for i, l in enumerate(self.loss):
+            if l['function'] is not None:
+                loss = l['function'](sr, hr)
+                effective_loss = l['weight'] * loss
+                losses.append(effective_loss)
+                self.log[-1, i] += effective_loss.item()
+            elif l['type'] == 'DIS':
+                self.log[-1, i] += self.loss[i - 1]['function'].loss
+
+        loss_sum = sum(losses)
+        if len(self.loss) > 1:
+            self.log[-1, -1] += loss_sum.item()
+
+        return loss_sum
+
+    def step(self):
+        for l in self.get_loss_module():
+            if hasattr(l, 'scheduler'):
+                l.scheduler.step()
+
+    def start_log(self):
+        self.log = torch.cat((self.log, torch.zeros(1, len(self.loss))))
+
+    def end_log(self, n_batches):
+        self.log[-1].div_(n_batches)
+
+    def display_loss(self, batch):
+        n_samples = batch + 1
+        log = []
+        for l, c in zip(self.loss, self.log[-1]):
+            log.append('[{}: {:.4f}]'.format(l['type'], c / n_samples))
+
+        return ''.join(log)
+
+    def plot_loss(self, apath, epoch):
+        axis = np.linspace(1, epoch, epoch)
+        for i, l in enumerate(self.loss):
+            label = '{} Loss'.format(l['type'])
+            fig = plt.figure()
+            plt.title(label)
+            plt.plot(axis, self.log[:, i].numpy(), label=label)
+            plt.legend()
+            plt.xlabel('Epochs')
+            plt.ylabel('Loss')
+            plt.grid(True)
+            plt.savefig(os.path.join(apath, 'loss_{}.pdf'.format(l['type'])))
+            plt.close(fig)
+
+    def get_loss_module(self):
+        if self.n_GPUs == 1:
+            return self.loss_module
+        else:
+            return self.loss_module.module
+
+    def save(self, apath):
+        torch.save(self.state_dict(), os.path.join(apath, 'loss.pt'))
+        torch.save(self.log, os.path.join(apath, 'loss_log.pt'))
+
+    def load(self, apath, cpu=False):
+        if cpu:
+            kwargs = {'map_location': lambda storage, loc: storage}
+        else:
+            kwargs = {}
+
+        self.load_state_dict(torch.load(
+            os.path.join(apath, 'loss.pt'),
+            **kwargs
+        ))
+        self.log = torch.load(os.path.join(apath, 'loss_log.pt'))
+        for l in self.get_loss_module():
+            if hasattr(l, 'scheduler'):
+                for _ in range(len(self.log)): l.scheduler.step()
+

src/loss/adversarial.py

+import utility
+from types import SimpleNamespace
+
+from model import common
+from loss import discriminator
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+
+class Adversarial(nn.Module):
+    def __init__(self, args, gan_type):
+        super(Adversarial, self).__init__()
+        self.gan_type = gan_type
+        self.gan_k = args.gan_k
+        self.dis = discriminator.Discriminator(args)
+        if gan_type == 'WGAN_GP':
+            # see https://arxiv.org/pdf/1704.00028.pdf pp.4
+            optim_dict = {
+                'optimizer': 'ADAM',
+                'betas': (0, 0.9),
+                'epsilon': 1e-8,
+                'lr': 1e-5,
+                'weight_decay': args.weight_decay,
+                'decay': args.decay,
+                'gamma': args.gamma
+            }
+            optim_args = SimpleNamespace(**optim_dict)
+        else:
+            optim_args = args
+
+        self.optimizer = utility.make_optimizer(optim_args, self.dis)
+
+    def forward(self, fake, real):
+        # updating discriminator...
+        self.loss = 0
+        fake_detach = fake.detach()     # do not backpropagate through G
+        for _ in range(self.gan_k):
+            self.optimizer.zero_grad()
+            # d: B x 1 tensor
+            d_fake = self.dis(fake_detach)
+            d_real = self.dis(real)
+            retain_graph = False
+            if self.gan_type == 'GAN':
+                loss_d = self.bce(d_real, d_fake)
+            elif self.gan_type.find('WGAN') >= 0:
+                loss_d = (d_fake - d_real).mean()
+                if self.gan_type.find('GP') >= 0:
+                    epsilon = torch.rand_like(fake).view(-1, 1, 1, 1)
+                    hat = fake_detach.mul(1 - epsilon) + real.mul(epsilon)
+                    hat.requires_grad = True
+                    d_hat = self.dis(hat)
+                    gradients = torch.autograd.grad(
+                        outputs=d_hat.sum(), inputs=hat,
+                        retain_graph=True, create_graph=True, only_inputs=True
+                    )[0]
+                    gradients = gradients.view(gradients.size(0), -1)
+                    gradient_norm = gradients.norm(2, dim=1)
+                    gradient_penalty = 10 * gradient_norm.sub(1).pow(2).mean()
+                    loss_d += gradient_penalty
+            # from ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks
+            elif self.gan_type == 'RGAN':
+                better_real = d_real - d_fake.mean(dim=0, keepdim=True)
+                better_fake = d_fake - d_real.mean(dim=0, keepdim=True)
+                loss_d = self.bce(better_real, better_fake)
+                retain_graph = True
+
+            # Discriminator update
+            self.loss += loss_d.item()
+            loss_d.backward(retain_graph=retain_graph)
+            self.optimizer.step()
+
+            if self.gan_type == 'WGAN':
+                for p in self.dis.parameters():
+                    p.data.clamp_(-1, 1)
+
+        self.loss /= self.gan_k
+
+        # updating generator...
+        d_fake_bp = self.dis(fake)      # for backpropagation, use fake as it is
+        if self.gan_type == 'GAN':
+            label_real = torch.ones_like(d_fake_bp)
+            loss_g = F.binary_cross_entropy_with_logits(d_fake_bp, label_real)
+        elif self.gan_type.find('WGAN') >= 0:
+            loss_g = -d_fake_bp.mean()
+        elif self.gan_type == 'RGAN':
+            better_real = d_real - d_fake_bp.mean(dim=0, keepdim=True)
+            better_fake = d_fake_bp - d_real.mean(dim=0, keepdim=True)
+            loss_g = self.bce(better_fake, better_real)
+
+        # Generator loss
+        return loss_g
+    
+    def state_dict(self, *args, **kwargs):
+        state_discriminator = self.dis.state_dict(*args, **kwargs)
+        state_optimizer = self.optimizer.state_dict()
+
+        return dict(**state_discriminator, **state_optimizer)
+
+    def bce(self, real, fake):
+        label_real = torch.ones_like(real)
+        label_fake = torch.zeros_like(fake)
+        bce_real = F.binary_cross_entropy_with_logits(real, label_real)
+        bce_fake = F.binary_cross_entropy_with_logits(fake, label_fake)
+        bce_loss = bce_real + bce_fake
+        return bce_loss
+               
+# Some references
+# https://github.com/kuc2477/pytorch-wgan-gp/blob/master/model.py
+# OR
+# https://github.com/caogang/wgan-gp/blob/master/gan_cifar10.py

src/loss/discriminator.py

+from model import common
+
+import torch.nn as nn
+
+class Discriminator(nn.Module):
+    '''
+        output is not normalized
+    '''
+    def __init__(self, args):
+        super(Discriminator, self).__init__()
+
+        in_channels = args.n_colors
+        out_channels = 64
+        depth = 7
+
+        def _block(_in_channels, _out_channels, stride=1):
+            return nn.Sequential(
+                nn.Conv2d(
+                    _in_channels,
+                    _out_channels,
+                    3,
+                    padding=1,
+                    stride=stride,
+                    bias=False
+                ),
+                nn.BatchNorm2d(_out_channels),
+                nn.LeakyReLU(negative_slope=0.2, inplace=True)
+            )
+
+        m_features = [_block(in_channels, out_channels)]
+        for i in range(depth):
+            in_channels = out_channels
+            if i % 2 == 1:
+                stride = 1
+                out_channels *= 2
+            else:
+                stride = 2
+            m_features.append(_block(in_channels, out_channels, stride=stride))
+
+        patch_size = args.patch_size // (2**((depth + 1) // 2))
+        m_classifier = [
+            nn.Linear(out_channels * patch_size**2, 1024),
+            nn.LeakyReLU(negative_slope=0.2, inplace=True),
+            nn.Linear(1024, 1)
+        ]
+
+        self.features = nn.Sequential(*m_features)
+        self.classifier = nn.Sequential(*m_classifier)
+
+    def forward(self, x):
+        features = self.features(x)
+        output = self.classifier(features.view(features.size(0), -1))
+
+        return output
+

src/loss/vgg.py

+from model import common
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+
+class VGG(nn.Module):
+    def __init__(self, conv_index, rgb_range=1):
+        super(VGG, self).__init__()
+        vgg_features = models.vgg19(pretrained=True).features
+        modules = [m for m in vgg_features]
+        if conv_index.find('22') >= 0:
+            self.vgg = nn.Sequential(*modules[:8])
+        elif conv_index.find('54') >= 0:
+            self.vgg = nn.Sequential(*modules[:35])
+
+        vgg_mean = (0.485, 0.456, 0.406)
+        vgg_std = (0.229 * rgb_range, 0.224 * rgb_range, 0.225 * rgb_range)
+        self.sub_mean = common.MeanShift(rgb_range, vgg_mean, vgg_std)
+        for p in self.parameters():
+            p.requires_grad = False
+
+    def forward(self, sr, hr):
+        def _forward(x):
+            x = self.sub_mean(x)
+            x = self.vgg(x)
+            return x
+            
+        vgg_sr = _forward(sr)
+        with torch.no_grad():
+            vgg_hr = _forward(hr.detach())
+
+        loss = F.mse_loss(vgg_sr, vgg_hr)
+
+        return loss

src/main.py

+import torch
+
+import utility
+import data
+import model
+import loss
+from option import args
+from trainer import Trainer
+from model.sspcab_torch import MY_SSPCAB
+
+torch.manual_seed(args.seed)
+checkpoint = utility.checkpoint(args)
+
+def main():
+    global model
+    if args.data_test == ['video']:
+        from videotester import VideoTester
+        model = model.Model(args, checkpoint)
+        t = VideoTester(args, model, checkpoint)
+        t.test()
+    else:
+        if checkpoint.ok:
+            loader = data.Data(args)
+            _model = model.Model(args, checkpoint)
+            _loss = loss.Loss(args, checkpoint) if not args.test_only else None
+            sub_model = MY_SSPCAB(channels=3, reduction_ratio=8)
+            sub_model.to('cuda')
+            t = Trainer(args, loader, _model, _loss, checkpoint, sub_model)
+            while not t.terminate():
+                t.train()
+                t.test()
+
+            checkpoint.done()
+
+if __name__ == '__main__':
+    main()

src/model/__init__.py

+import os
+from importlib import import_module
+
+import torch
+import torch.nn as nn
+import torch.nn.parallel as P
+import torch.utils.model_zoo
+
+class Model(nn.Module):
+    def __init__(self, args, ckp):
+        super(Model, self).__init__()
+        print('Making model...')
+
+        self.scale = args.scale
+        self.idx_scale = 0
+        self.input_large = (args.model == 'VDSR')
+        self.self_ensemble = args.self_ensemble
+        self.chop = args.chop
+        self.precision = args.precision
+        self.cpu = args.cpu
+        if self.cpu:
+            self.device = torch.device('cpu')
+        else:
+            #if torch.backends.mps.is_available():
+            #    self.device = torch.device('mps')
+            if torch.cuda.is_available():
+                self.device = torch.device('cuda')
+            else:
+                self.device = torch.device('cpu')
+
+        self.n_GPUs = args.n_GPUs
+        self.save_models = args.save_models
+
+        module = import_module('model.' + args.model.lower())
+        self.model = module.make_model(args).to(self.device)
+        if args.precision == 'half':
+            self.model.half()
+
+        self.load(
+            ckp.get_path('model'),
+            pre_train=args.pre_train,
+            resume=args.resume,
+            cpu=args.cpu
+        )
+        print(self.model, file=ckp.log_file)
+
+    def forward(self, x, idx_scale):
+        self.idx_scale = idx_scale
+        if hasattr(self.model, 'set_scale'):
+            self.model.set_scale(idx_scale)
+
+        if self.training:
+            if self.n_GPUs > 1:
+                return P.data_parallel(self.model, x, range(self.n_GPUs))
+            else:
+                return self.model(x)
+        else:
+            if self.chop:
+                forward_function = self.forward_chop
+            else:
+                forward_function = self.model.forward
+
+            if self.self_ensemble:
+                return self.forward_x8(x, forward_function=forward_function)
+            else:
+                return forward_function(x)
+
+    def save(self, apath, epoch, is_best=False):
+        save_dirs = [os.path.join(apath, 'model_latest.pt')]
+
+        if is_best:
+            save_dirs.append(os.path.join(apath, 'model_best.pt'))
+        if self.save_models:
+            save_dirs.append(
+                os.path.join(apath, 'model_{}.pt'.format(epoch))
+            )
+
+        for s in save_dirs:
+            torch.save(self.model.state_dict(), s)
+
+    def load(self, apath, pre_train='', resume=-1, cpu=False):
+        load_from = None
+        kwargs = {}
+        if cpu:
+            kwargs = {'map_location': lambda storage, loc: storage}
+        else:
+            kwargs = {'map_location': self.device}
+
+        if resume == -1:
+            load_from = torch.load(
+                os.path.join(apath, 'model_latest.pt'),
+                **kwargs
+            )
+        elif resume == 0:
+            if pre_train == 'download':
+                print('Download the model')
+                dir_model = os.path.join('..', 'models')
+                os.makedirs(dir_model, exist_ok=True)
+                load_from = torch.utils.model_zoo.load_url(
+                    self.model.url,
+                    model_dir=dir_model,
+                    **kwargs
+                )
+            elif pre_train:
+                print('Load the model from {}'.format(pre_train))
+                load_from = torch.load(pre_train, **kwargs)
+        else:
+            load_from = torch.load(
+                os.path.join(apath, 'model_{}.pt'.format(resume)),
+                **kwargs
+            )
+
+        if load_from:
+            self.model.load_state_dict(load_from, strict=False)
+
+    def forward_chop(self, *args, shave=10, min_size=160000):
+        scale = 1 if self.input_large else self.scale[self.idx_scale]
+        n_GPUs = min(self.n_GPUs, 4)
+        # height, width
+        h, w = args[0].size()[-2:]
+
+        top = slice(0, h//2 + shave)
+        bottom = slice(h - h//2 - shave, h)
+        left = slice(0, w//2 + shave)
+        right = slice(w - w//2 - shave, w)
+        x_chops = [torch.cat([
+            a[..., top, left],
+            a[..., top, right],
+            a[..., bottom, left],
+            a[..., bottom, right]
+        ]) for a in args]
+
+        y_chops = []
+        if h * w < 4 * min_size:
+            for i in range(0, 4, n_GPUs):
+                x = [x_chop[i:(i + n_GPUs)] for x_chop in x_chops]
+                y = P.data_parallel(self.model, *x, range(n_GPUs))
+                if not isinstance(y, list): y = [y]
+                if not y_chops:
+                    y_chops = [[c for c in _y.chunk(n_GPUs, dim=0)] for _y in y]
+                else:
+                    for y_chop, _y in zip(y_chops, y):
+                        y_chop.extend(_y.chunk(n_GPUs, dim=0))
+        else:
+            for p in zip(*x_chops):
+                y = self.forward_chop(*p, shave=shave, min_size=min_size)
+                if not isinstance(y, list): y = [y]
+                if not y_chops:
+                    y_chops = [[_y] for _y in y]
+                else:
+                    for y_chop, _y in zip(y_chops, y): y_chop.append(_y)
+
+        h *= scale
+        w *= scale
+        top = slice(0, h//2)
+        bottom = slice(h - h//2, h)
+        bottom_r = slice(h//2 - h, None)
+        left = slice(0, w//2)
+        right = slice(w - w//2, w)
+        right_r = slice(w//2 - w, None)
+
+        # batch size, number of color channels
+        b, c = y_chops[0][0].size()[:-2]
+        y = [y_chop[0].new(b, c, h, w) for y_chop in y_chops]
+        for y_chop, _y in zip(y_chops, y):
+            _y[..., top, left] = y_chop[0][..., top, left]
+            _y[..., top, right] = y_chop[1][..., top, right_r]
+            _y[..., bottom, left] = y_chop[2][..., bottom_r, left]
+            _y[..., bottom, right] = y_chop[3][..., bottom_r, right_r]
+
+        if len(y) == 1: y = y[0]
+
+        return y
+
+    def forward_x8(self, *args, forward_function=None):
+        def _transform(v, op):
+            if self.precision != 'single': v = v.float()
+
+            v2np = v.data.cpu().numpy()
+            if op == 'v':
+                tfnp = v2np[:, :, :, ::-1].copy()
+            elif op == 'h':
+                tfnp = v2np[:, :, ::-1, :].copy()
+            elif op == 't':
+                tfnp = v2np.transpose((0, 1, 3, 2)).copy()
+
+            ret = torch.Tensor(tfnp).to(self.device)
+            if self.precision == 'half': ret = ret.half()
+
+            return ret
+
+        list_x = []
+        for a in args:
+            x = [a]
+            for tf in 'v', 'h', 't': x.extend([_transform(_x, tf) for _x in x])
+
+            list_x.append(x)
+
+        list_y = []
+        for x in zip(*list_x):
+            y = forward_function(*x)
+            if not isinstance(y, list): y = [y]
+            if not list_y:
+                list_y = [[_y] for _y in y]
+            else:
+                for _list_y, _y in zip(list_y, y): _list_y.append(_y)
+
+        for _list_y in list_y:
+            for i in range(len(_list_y)):
+                if i > 3:
+                    _list_y[i] = _transform(_list_y[i], 't')
+                if i % 4 > 1:
+                    _list_y[i] = _transform(_list_y[i], 'h')
+                if (i % 4) % 2 == 1:
+                    _list_y[i] = _transform(_list_y[i], 'v')
+
+        y = [torch.cat(_y, dim=0).mean(dim=0, keepdim=True) for _y in list_y]
+        if len(y) == 1: y = y[0]
+
+        return y

src/model/common.py

+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def default_conv(in_channels, out_channels, kernel_size, bias=True):
+    return nn.Conv2d(
+        in_channels, out_channels, kernel_size,
+        padding=(kernel_size//2), bias=bias)
+
+class MeanShift(nn.Conv2d):
+    def __init__(
+        self, rgb_range,
+        rgb_mean=(0.4488, 0.4371, 0.4040), rgb_std=(1.0, 1.0, 1.0), sign=-1):
+
+        super(MeanShift, self).__init__(3, 3, kernel_size=1)
+        std = torch.Tensor(rgb_std)
+        self.weight.data = torch.eye(3).view(3, 3, 1, 1) / std.view(3, 1, 1, 1)
+        self.bias.data = sign * rgb_range * torch.Tensor(rgb_mean) / std
+        for p in self.parameters():
+            p.requires_grad = False
+
+class BasicBlock(nn.Sequential):
+    def __init__(
+        self, conv, in_channels, out_channels, kernel_size, stride=1, bias=False,
+        bn=True, act=nn.ReLU(True)):
+
+        m = [conv(in_channels, out_channels, kernel_size, bias=bias)]
+        if bn:
+            m.append(nn.BatchNorm2d(out_channels))
+        if act is not None:
+            m.append(act)
+
+        super(BasicBlock, self).__init__(*m)
+
+class ResBlock(nn.Module):
+    def __init__(
+        self, conv, n_feats, kernel_size,
+        bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
+
+        super(ResBlock, self).__init__()
+        m = []
+        for i in range(2):
+            m.append(conv(n_feats, n_feats, kernel_size, bias=bias))
+            if bn:
+                m.append(nn.BatchNorm2d(n_feats))
+            if i == 0:
+                m.append(act)
+
+        self.body = nn.Sequential(*m)
+        self.res_scale = res_scale
+
+    def forward(self, x):
+        res = self.body(x).mul(self.res_scale)
+        res += x
+
+        return res
+
+class Upsampler(nn.Sequential):
+    def __init__(self, conv, scale, n_feats, bn=False, act=False, bias=True):
+
+        m = []
+        if (scale & (scale - 1)) == 0:    # Is scale = 2^n?
+            for _ in range(int(math.log(scale, 2))):
+                m.append(conv(n_feats, 4 * n_feats, 3, bias))
+                m.append(nn.PixelShuffle(2))
+                if bn:
+                    m.append(nn.BatchNorm2d(n_feats))
+                if act == 'relu':
+                    m.append(nn.ReLU(True))
+                elif act == 'prelu':
+                    m.append(nn.PReLU(n_feats))
+
+        elif scale == 3:
+            m.append(conv(n_feats, 9 * n_feats, 3, bias))
+            m.append(nn.PixelShuffle(3))
+            if bn:
+                m.append(nn.BatchNorm2d(n_feats))
+            if act == 'relu':
+                m.append(nn.ReLU(True))
+            elif act == 'prelu':
+                m.append(nn.PReLU(n_feats))
+        else:
+            raise NotImplementedError
+
+        super(Upsampler, self).__init__(*m)
+

src/model/ddbpn.py

+# Deep Back-Projection Networks For Super-Resolution
+# https://arxiv.org/abs/1803.02735
+
+from model import common
+
+import torch
+import torch.nn as nn
+
+
+def make_model(args, parent=False):
+    return DDBPN(args)
+
+def projection_conv(in_channels, out_channels, scale, up=True):
+    kernel_size, stride, padding = {
+        2: (6, 2, 2),
+        4: (8, 4, 2),
+        8: (12, 8, 2)
+    }[scale]
+    if up:
+        conv_f = nn.ConvTranspose2d
+    else:
+        conv_f = nn.Conv2d
+
+    return conv_f(
+        in_channels, out_channels, kernel_size,
+        stride=stride, padding=padding
+    )
+
+class DenseProjection(nn.Module):
+    def __init__(self, in_channels, nr, scale, up=True, bottleneck=True):
+        super(DenseProjection, self).__init__()
+        if bottleneck:
+            self.bottleneck = nn.Sequential(*[
+                nn.Conv2d(in_channels, nr, 1),
+                nn.PReLU(nr)
+            ])
+            inter_channels = nr
+        else:
+            self.bottleneck = None
+            inter_channels = in_channels
+
+        self.conv_1 = nn.Sequential(*[
+            projection_conv(inter_channels, nr, scale, up),
+            nn.PReLU(nr)
+        ])
+        self.conv_2 = nn.Sequential(*[
+            projection_conv(nr, inter_channels, scale, not up),
+            nn.PReLU(inter_channels)
+        ])
+        self.conv_3 = nn.Sequential(*[
+            projection_conv(inter_channels, nr, scale, up),
+            nn.PReLU(nr)
+        ])
+
+    def forward(self, x):
+        if self.bottleneck is not None:
+            x = self.bottleneck(x)
+
+        a_0 = self.conv_1(x)
+        b_0 = self.conv_2(a_0)
+        e = b_0.sub(x)
+        a_1 = self.conv_3(e)
+
+        out = a_0.add(a_1)
+
+        return out
+
+class DDBPN(nn.Module):
+    def __init__(self, args):
+        super(DDBPN, self).__init__()
+        scale = args.scale[0]
+
+        n0 = 128
+        nr = 32
+        self.depth = 6
+
+        rgb_mean = (0.4488, 0.4371, 0.4040)
+        rgb_std = (1.0, 1.0, 1.0)
+        self.sub_mean = common.MeanShift(args.rgb_range, rgb_mean, rgb_std)
+        initial = [
+            nn.Conv2d(args.n_colors, n0, 3, padding=1),
+            nn.PReLU(n0),
+            nn.Conv2d(n0, nr, 1),
+            nn.PReLU(nr)
+        ]
+        self.initial = nn.Sequential(*initial)
+
+        self.upmodules = nn.ModuleList()
+        self.downmodules = nn.ModuleList()
+        channels = nr
+        for i in range(self.depth):
+            self.upmodules.append(
+                DenseProjection(channels, nr, scale, True, i > 1)
+            )
+            if i != 0:
+                channels += nr
+        
+        channels = nr
+        for i in range(self.depth - 1):
+            self.downmodules.append(
+                DenseProjection(channels, nr, scale, False, i != 0)
+            )
+            channels += nr
+
+        reconstruction = [
+            nn.Conv2d(self.depth * nr, args.n_colors, 3, padding=1) 
+        ]
+        self.reconstruction = nn.Sequential(*reconstruction)
+
+        self.add_mean = common.MeanShift(args.rgb_range, rgb_mean, rgb_std, 1)
+
+    def forward(self, x):
+        x = self.sub_mean(x)
+        x = self.initial(x)
+
+        h_list = []
+        l_list = []
+        for i in range(self.depth - 1):
+            if i == 0:
+                l = x
+            else:
+                l = torch.cat(l_list, dim=1)
+            h_list.append(self.upmodules[i](l))
+            l_list.append(self.downmodules[i](torch.cat(h_list, dim=1)))
+        
+        h_list.append(self.upmodules[-1](torch.cat(l_list, dim=1)))
+        out = self.reconstruction(torch.cat(h_list, dim=1))
+        out = self.add_mean(out)
+
+        return out
+

src/model/edsr.py

+from model import common
+
+import torch.nn as nn
+
+url = {
+    'r16f64x2': 'https://cv.snu.ac.kr/research/EDSR/models/edsr_baseline_x2-1bc95232.pt',
+    'r16f64x3': 'https://cv.snu.ac.kr/research/EDSR/models/edsr_baseline_x3-abf2a44e.pt',
+    'r16f64x4': 'https://cv.snu.ac.kr/research/EDSR/models/edsr_baseline_x4-6b446fab.pt',
+    'r32f256x2': 'https://cv.snu.ac.kr/research/EDSR/models/edsr_x2-0edfb8a3.pt',
+    'r32f256x3': 'https://cv.snu.ac.kr/research/EDSR/models/edsr_x3-ea3ef2c6.pt',
+    'r32f256x4': 'https://cv.snu.ac.kr/research/EDSR/models/edsr_x4-4f62e9ef.pt'
+}
+
+def make_model(args, parent=False):
+    return EDSR(args)
+
+class EDSR(nn.Module):
+    def __init__(self, args, conv=common.default_conv):
+        super(EDSR, self).__init__()
+
+        n_resblocks = args.n_resblocks
+        n_feats = args.n_feats
+        kernel_size = 3 
+        scale = args.scale[0]
+        act = nn.ReLU(True)
+        url_name = 'r{}f{}x{}'.format(n_resblocks, n_feats, scale)
+        if url_name in url:
+            self.url = url[url_name]
+        else:
+            self.url = None
+        self.sub_mean = common.MeanShift(args.rgb_range)
+        self.add_mean = common.MeanShift(args.rgb_range, sign=1)
+
+        # define head module
+        m_head = [conv(args.n_colors, n_feats, kernel_size)]
+
+        # define body module
+        m_body = [
+            common.ResBlock(
+                conv, n_feats, kernel_size, act=act, res_scale=args.res_scale
+            ) for _ in range(n_resblocks)
+        ]
+        m_body.append(conv(n_feats, n_feats, kernel_size))
+
+        # define tail module
+        m_tail = [
+            common.Upsampler(conv, scale, n_feats, act=False),
+            conv(n_feats, args.n_colors, kernel_size)
+        ]
+
+        self.head = nn.Sequential(*m_head)
+        self.body = nn.Sequential(*m_body)
+        self.tail = nn.Sequential(*m_tail)
+
+    def forward(self, x):
+        x = self.sub_mean(x)
+        x = self.head(x)
+
+        res = self.body(x)
+        res += x
+
+        x = self.tail(res)
+        x = self.add_mean(x)
+
+        return x 
+
+    def load_state_dict(self, state_dict, strict=True):
+        own_state = self.state_dict()
+        for name, param in state_dict.items():
+            if name in own_state:
+                if isinstance(param, nn.Parameter):
+                    param = param.data
+                try:
+                    own_state[name].copy_(param)
+                except Exception:
+                    if name.find('tail') == -1:
+                        raise RuntimeError('While copying the parameter named {}, '
+                                           'whose dimensions in the model are {} and '
+                                           'whose dimensions in the checkpoint are {}.'
+                                           .format(name, own_state[name].size(), param.size()))
+            elif strict:
+                if name.find('tail') == -1:
+                    raise KeyError('unexpected key "{}" in state_dict'
+                                   .format(name))
+

src/model/mdsr.py

+from model import common
+
+import torch.nn as nn
+
+url = {
+    'r16f64': 'https://cv.snu.ac.kr/research/EDSR/models/mdsr_baseline-a00cab12.pt',
+    'r80f64': 'https://cv.snu.ac.kr/research/EDSR/models/mdsr-4a78bedf.pt'
+}
+
+def make_model(args, parent=False):
+    return MDSR(args)
+
+class MDSR(nn.Module):
+    def __init__(self, args, conv=common.default_conv):
+        super(MDSR, self).__init__()
+        n_resblocks = args.n_resblocks
+        n_feats = args.n_feats
+        kernel_size = 3
+        act = nn.ReLU(True)
+        self.scale_idx = 0
+        self.url = url['r{}f{}'.format(n_resblocks, n_feats)]
+        self.sub_mean = common.MeanShift(args.rgb_range)
+        self.add_mean = common.MeanShift(args.rgb_range, sign=1)
+
+        m_head = [conv(args.n_colors, n_feats, kernel_size)]
+
+        self.pre_process = nn.ModuleList([
+            nn.Sequential(
+                common.ResBlock(conv, n_feats, 5, act=act),
+                common.ResBlock(conv, n_feats, 5, act=act)
+            ) for _ in args.scale
+        ])
+
+        m_body = [
+            common.ResBlock(
+                conv, n_feats, kernel_size, act=act
+            ) for _ in range(n_resblocks)
+        ]
+        m_body.append(conv(n_feats, n_feats, kernel_size))
+
+        self.upsample = nn.ModuleList([
+            common.Upsampler(conv, s, n_feats, act=False) for s in args.scale
+        ])
+
+        m_tail = [conv(n_feats, args.n_colors, kernel_size)]
+
+        self.head = nn.Sequential(*m_head)
+        self.body = nn.Sequential(*m_body)
+        self.tail = nn.Sequential(*m_tail)
+
+    def forward(self, x):
+        x = self.sub_mean(x)
+        x = self.head(x)
+        x = self.pre_process[self.scale_idx](x)
+
+        res = self.body(x)
+        res += x
+
+        x = self.upsample[self.scale_idx](res)
+        x = self.tail(x)
+        x = self.add_mean(x)
+
+        return x
+
+    def set_scale(self, scale_idx):
+        self.scale_idx = scale_idx
+

src/model/rcan.py

+## ECCV-2018-Image Super-Resolution Using Very Deep Residual Channel Attention Networks
+## https://arxiv.org/abs/1807.02758
+from model import common
+
+import torch.nn as nn
+
+def make_model(args, parent=False):
+    return RCAN(args)
+
+## Channel Attention (CA) Layer
+class CALayer(nn.Module):
+    def __init__(self, channel, reduction=16):
+        super(CALayer, self).__init__()
+        # global average pooling: feature --> point
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        # feature channel downscale and upscale --> channel weight
+        self.conv_du = nn.Sequential(
+                nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=True),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=True),
+                nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        y = self.avg_pool(x)
+        y = self.conv_du(y)
+        return x * y
+
+## Residual Channel Attention Block (RCAB)
+class RCAB(nn.Module):
+    def __init__(
+        self, conv, n_feat, kernel_size, reduction,
+        bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
+
+        super(RCAB, self).__init__()
+        modules_body = []
+        for i in range(2):
+            modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
+            if bn: modules_body.append(nn.BatchNorm2d(n_feat))
+            if i == 0: modules_body.append(act)
+        modules_body.append(CALayer(n_feat, reduction))
+        self.body = nn.Sequential(*modules_body)
+        self.res_scale = res_scale
+
+    def forward(self, x):
+        res = self.body(x)
+        #res = self.body(x).mul(self.res_scale)
+        res += x
+        return res
+
+## Residual Group (RG)
+class ResidualGroup(nn.Module):
+    def __init__(self, conv, n_feat, kernel_size, reduction, act, res_scale, n_resblocks):
+        super(ResidualGroup, self).__init__()
+        modules_body = []
+        modules_body = [
+            RCAB(
+                conv, n_feat, kernel_size, reduction, bias=True, bn=False, act=nn.ReLU(True), res_scale=1) \
+            for _ in range(n_resblocks)]
+        modules_body.append(conv(n_feat, n_feat, kernel_size))
+        self.body = nn.Sequential(*modules_body)
+
+    def forward(self, x):
+        res = self.body(x)
+        res += x
+        return res
+
+## Residual Channel Attention Network (RCAN)
+class RCAN(nn.Module):
+    def __init__(self, args, conv=common.default_conv):
+        super(RCAN, self).__init__()
+        
+        n_resgroups = args.n_resgroups
+        n_resblocks = args.n_resblocks
+        n_feats = args.n_feats
+        kernel_size = 3
+        reduction = args.reduction 
+        scale = args.scale[0]
+        act = nn.ReLU(True)
+        
+        # RGB mean for DIV2K
+        self.sub_mean = common.MeanShift(args.rgb_range)
+        
+        # define head module
+        modules_head = [conv(args.n_colors, n_feats, kernel_size)]
+
+        # define body module
+        modules_body = [
+            ResidualGroup(
+                conv, n_feats, kernel_size, reduction, act=act, res_scale=args.res_scale, n_resblocks=n_resblocks) \
+            for _ in range(n_resgroups)]
+
+        modules_body.append(conv(n_feats, n_feats, kernel_size))
+
+        # define tail module
+        modules_tail = [
+            common.Upsampler(conv, scale, n_feats, act=False),
+            conv(n_feats, args.n_colors, kernel_size)]
+
+        self.add_mean = common.MeanShift(args.rgb_range, sign=1)
+
+        self.head = nn.Sequential(*modules_head)
+        self.body = nn.Sequential(*modules_body)
+        self.tail = nn.Sequential(*modules_tail)
+
+    def forward(self, x):
+        x = self.sub_mean(x)
+        x = self.head(x)
+
+        res = self.body(x)
+        res += x
+
+        x = self.tail(res)
+        x = self.add_mean(x)
+
+        return x 
+
+    def load_state_dict(self, state_dict, strict=False):
+        own_state = self.state_dict()
+        for name, param in state_dict.items():
+            if name in own_state:
+                if isinstance(param, nn.Parameter):
+                    param = param.data
+                try:
+                    own_state[name].copy_(param)
+                except Exception:
+                    if name.find('tail') >= 0:
+                        print('Replace pre-trained upsampler to new one...')
+                    else:
+                        raise RuntimeError('While copying the parameter named {}, '
+                                           'whose dimensions in the model are {} and '
+                                           'whose dimensions in the checkpoint are {}.'
+                                           .format(name, own_state[name].size(), param.size()))
+            elif strict:
+                if name.find('tail') == -1:
+                    raise KeyError('unexpected key "{}" in state_dict'
+                                   .format(name))
+
+        if strict:
+            missing = set(own_state.keys()) - set(state_dict.keys())
+            if len(missing) > 0:
+                raise KeyError('missing keys in state_dict: "{}"'.format(missing))

src/model/rdn.py

+# Residual Dense Network for Image Super-Resolution
+# https://arxiv.org/abs/1802.08797
+
+from model import common
+
+import torch
+import torch.nn as nn
+
+
+def make_model(args, parent=False):
+    return RDN(args)
+
+class RDB_Conv(nn.Module):
+    def __init__(self, inChannels, growRate, kSize=3):
+        super(RDB_Conv, self).__init__()
+        Cin = inChannels
+        G  = growRate
+        self.conv = nn.Sequential(*[
+            nn.Conv2d(Cin, G, kSize, padding=(kSize-1)//2, stride=1),
+            nn.ReLU()
+        ])
+
+    def forward(self, x):
+        out = self.conv(x)
+        return torch.cat((x, out), 1)
+
+class RDB(nn.Module):
+    def __init__(self, growRate0, growRate, nConvLayers, kSize=3):
+        super(RDB, self).__init__()
+        G0 = growRate0
+        G  = growRate
+        C  = nConvLayers
+        
+        convs = []
+        for c in range(C):
+            convs.append(RDB_Conv(G0 + c*G, G))
+        self.convs = nn.Sequential(*convs)
+        
+        # Local Feature Fusion
+        self.LFF = nn.Conv2d(G0 + C*G, G0, 1, padding=0, stride=1)
+
+    def forward(self, x):
+        return self.LFF(self.convs(x)) + x
+
+class RDN(nn.Module):
+    def __init__(self, args):
+        super(RDN, self).__init__()
+        r = args.scale[0]
+        G0 = args.G0
+        kSize = args.RDNkSize
+
+        # number of RDB blocks, conv layers, out channels
+        self.D, C, G = {
+            'A': (20, 6, 32),
+            'B': (16, 8, 64),
+        }[args.RDNconfig]
+
+        # Shallow feature extraction net
+        self.SFENet1 = nn.Conv2d(args.n_colors, G0, kSize, padding=(kSize-1)//2, stride=1)
+        self.SFENet2 = nn.Conv2d(G0, G0, kSize, padding=(kSize-1)//2, stride=1)
+
+        # Redidual dense blocks and dense feature fusion
+        self.RDBs = nn.ModuleList()
+        for i in range(self.D):
+            self.RDBs.append(
+                RDB(growRate0 = G0, growRate = G, nConvLayers = C)
+            )
+
+        # Global Feature Fusion
+        self.GFF = nn.Sequential(*[
+            nn.Conv2d(self.D * G0, G0, 1, padding=0, stride=1),
+            nn.Conv2d(G0, G0, kSize, padding=(kSize-1)//2, stride=1)
+        ])
+
+        # Up-sampling net
+        if r == 2 or r == 3:
+            self.UPNet = nn.Sequential(*[
+                nn.Conv2d(G0, G * r * r, kSize, padding=(kSize-1)//2, stride=1),
+                nn.PixelShuffle(r),
+                nn.Conv2d(G, args.n_colors, kSize, padding=(kSize-1)//2, stride=1)
+            ])
+        elif r == 4:
+            self.UPNet = nn.Sequential(*[
+                nn.Conv2d(G0, G * 4, kSize, padding=(kSize-1)//2, stride=1),
+                nn.PixelShuffle(2),
+                nn.Conv2d(G, G * 4, kSize, padding=(kSize-1)//2, stride=1),
+                nn.PixelShuffle(2),
+                nn.Conv2d(G, args.n_colors, kSize, padding=(kSize-1)//2, stride=1)
+            ])
+        else:
+            raise ValueError("scale must be 2 or 3 or 4.")
+
+    def forward(self, x):
+        f__1 = self.SFENet1(x)
+        x  = self.SFENet2(f__1)
+
+        RDBs_out = []
+        for i in range(self.D):
+            x = self.RDBs[i](x)
+            RDBs_out.append(x)
+
+        x = self.GFF(torch.cat(RDBs_out,1))
+        x += f__1
+
+        return self.UPNet(x)