Upload folder using huggingface_hub

CAR/README.md

@@ -0,0 +1,103 @@
+# Pyramid Attention for Image Restoration
+This repository is for PANet and PA-EDSR introduced in the following paper
+
+[Yiqun Mei](http://yiqunm2.web.illinois.edu/), [Yuchen Fan](https://scholar.google.com/citations?user=BlfdYL0AAAAJ&hl=en), [Yulun Zhang](http://yulunzhang.com/), [Jiahui Yu](https://jiahuiyu.com/), [Yuqian Zhou](https://yzhouas.github.io/), [Ding Liu](https://scholar.google.com/citations?user=PGtHUI0AAAAJ&hl=en), [Yun Fu](http://www1.ece.neu.edu/~yunfu/), [Thomas S. Huang](http://ifp-uiuc.github.io/) and [Honghui Shi](https://www.humphreyshi.com/) "Pyramid Attention for Image Restoration", [[Arxiv]](https://arxiv.org/abs/2004.13824) 
+
+The code is built on [EDSR (PyTorch)](https://github.com/thstkdgus35/EDSR-PyTorch) & [RNAN](https://github.com/yulunzhang/RNAN) and tested on Ubuntu 18.04 environment (Python3.6, PyTorch_1.1) with Titan X/1080Ti/V100 GPUs.
+
+## Contents
+1. [Train](#train)
+2. [Test](#test)
+3. [Results](#results)
+4. [Citation](#citation)
+5. [Acknowledgements](#acknowledgements)
+
+## Train
+### Prepare training data 
+
+1. Download DIV2K training data (800 training + 100 validtion images) from [DIV2K dataset](https://data.vision.ee.ethz.ch/cvl/DIV2K/) or [SNU_CVLab](https://cv.snu.ac.kr/research/EDSR/DIV2K.tar).
+
+2. Specify '--dir_data' based on the HR and LR images path.
+
+3. Organize training data like:
+```bash
+DIV2K/
+├── DIV2K_train_HR
+├── DIV2K_train_LR_bicubic
+│   └── X10
+│   └── X20
+│   └── X30
+│   └── X40
+├── DIV2K_valid_HR
+└── DIV2K_valid_LR_bicubic
+    └── X10
+    └── X20
+    └── X30
+    └── X40
+```
+For more informaiton, please refer to [EDSR(PyTorch)](https://github.com/thstkdgus35/EDSR-PyTorch).
+
+### Begin to train
+
+1. (optional) All the pretrained models and visual results can be downloaded from [Google Drive](https://drive.google.com/open?id=1q9iUzqYX0fVRzDu4J6fvSPRosgOZoJJE).
+
+2. Cd to 'PANet-PyTorch/[Task]/code', run the following scripts to train models.
+
+    **You can use scripts in file 'demo.sb' to train and test models for our paper.**
+
+    ```bash
+    # Example Usage: Q=10 
+    python main.py --n_GPUs 2 --batch_size 16 --lr 1e-4 --decay 200-400-600-800 --save_models --n_resblocks 80 --model PANET --scale 10 --patch_size 48 --save PANET_Q10 --n_feats 64 --data_train DIV2K --chop
+    ```
+## Test
+### Quick start
+
+1. Cd to 'PANet-PyTorch/[Task]/code', run the following scripts.
+
+    **You can use scripts in file 'demo.sb' to produce results for our paper.**
+
+    ```bash
+    # No self-ensemble, use different testsets (Classic5, LIVE1) to reproduce the results in the paper.
+    # Example Usage: Q=40
+    python main.py --model PANET --save_results --n_GPUs 1 --chop --data_test classic5+LIVE1 --scale 40 --n_resblocks 80 --n_feats 64 --pre_train ../Q40.pt --test_only
+
+    ```
+
+### The whole test pipeline
+1. Prepare test data. Organize training data like:
+```bash
+benchmark/
+├── testset1
+│   └── HR 
+│   └── LR_bicubic
+│   	└── X10
+│   	└── ..
+├── testset2
+```
+
+2. Conduct image CAR. 
+
+    See **Quick start**
+3. Evaluate the results.
+
+    Run 'Evaluate_PSNR_SSIM.m' to obtain PSNR/SSIM values for paper.
+
+## Citation
+If you find the code helpful in your resarch or work, please cite the following papers.
+```
+@article{mei2020pyramid,
+  title={Pyramid Attention Networks for Image Restoration},
+  author={Mei, Yiqun and Fan, Yuchen and Zhang, Yulun and Yu, Jiahui and Zhou, Yuqian and Liu, Ding and Fu, Yun and Huang, Thomas S and Shi, Honghui},
+  journal={arXiv preprint arXiv:2004.13824},
+  year={2020}
+}
+@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}
+}
+```
+## Acknowledgements
+This code is built on [EDSR (PyTorch)](https://github.com/thstkdgus35/EDSR-PyTorch), [RNAN](https://github.com/yulunzhang/RNAN) and [generative-inpainting-pytorch](https://github.com/daa233/generative-inpainting-pytorch). We thank the authors for sharing their codes.

CAR/code/LICENSE

@@ -0,0 +1,21 @@
+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.

CAR/code/data/__init__.py

@@ -0,0 +1,52 @@
+from importlib import import_module
+#from dataloader import MSDataLoader
+from torch.utils.data import dataloader
+from torch.utils.data import ConcatDataset
+
+# This is a simple wrapper function for ConcatDataset
+class MyConcatDataset(ConcatDataset):
+    def __init__(self, datasets):
+        super(MyConcatDataset, self).__init__(datasets)
+        self.train = datasets[0].train
+
+    def set_scale(self, idx_scale):
+        for d in self.datasets:
+            if hasattr(d, 'set_scale'): d.set_scale(idx_scale)
+
+class Data:
+    def __init__(self, args):
+        self.loader_train = None
+        if not args.test_only:
+            datasets = []
+            for d in args.data_train:
+                module_name = d if d.find('DIV2K-Q') < 0 else 'DIV2KJPEG'
+                m = import_module('data.' + module_name.lower())
+                datasets.append(getattr(m, module_name)(args, name=d))
+
+            self.loader_train = dataloader.DataLoader(
+                MyConcatDataset(datasets),
+                batch_size=args.batch_size,
+                shuffle=True,
+                pin_memory=not args.cpu,
+                num_workers=args.n_threads,
+            )
+
+        self.loader_test = []
+        for d in args.data_test:
+            if d in ['CBSD68','classic5','LIVE1','Kodak24','Set5', 'Set14', 'B100', 'Urban100']:
+                m = import_module('data.benchmark')
+                testset = getattr(m, 'Benchmark')(args, train=False, name=d)
+            else:
+                module_name = d if d.find('DIV2K-Q') < 0 else 'DIV2KJPEG'
+                m = import_module('data.' + module_name.lower())
+                testset = getattr(m, module_name)(args, train=False, name=d)
+
+            self.loader_test.append(
+                dataloader.DataLoader(
+                    testset,
+                    batch_size=1,
+                    shuffle=False,
+                    pin_memory=not args.cpu,
+                    num_workers=args.n_threads,
+                )
+            )

CAR/code/data/benchmark.py

@@ -0,0 +1,25 @@
+import os
+
+from data import common
+from data import srdata
+
+import numpy as np
+
+import torch
+import torch.utils.data as data
+
+class Benchmark(srdata.SRData):
+    def __init__(self, args, name='', train=True, benchmark=True):
+        super(Benchmark, self).__init__(
+            args, name=name, train=train, benchmark=True
+        )
+
+    def _set_filesystem(self, dir_data):
+        self.apath = os.path.join(dir_data, 'benchmark', self.name)
+        self.dir_hr = os.path.join(self.apath, 'HR')
+        if self.input_large:
+            self.dir_lr = os.path.join(self.apath, 'LR_bicubicL')
+        else:
+            self.dir_lr = os.path.join(self.apath, 'LR_bicubic')
+        self.ext = ('','.jpg')
+

CAR/code/data/common.py

@@ -0,0 +1,79 @@
+import random
+
+import numpy as np
+import skimage.color as sc
+
+import torch
+
+def get_patch(*args, patch_size=96, scale=1, multi=False, input_large=False):
+    ih, iw = args[0].shape[:2]
+    # print('heelo')
+    # print(args[0].shape)
+
+    if not input_large:
+        p = 1 if multi else 1
+        tp = p * patch_size
+        ip = tp // 1
+    else:
+        tp = patch_size
+        ip = patch_size
+
+    ix = random.randrange(0, iw - ip + 1)
+    iy = random.randrange(0, ih - ip + 1)
+
+    if not input_large:
+        tx, ty = 1 * ix, 1 * iy
+    else:
+        tx, ty = ix, iy
+
+    ret = [
+        args[0][iy:iy + ip, ix:ix + ip],
+        *[a[ty:ty + tp, tx:tx + tp] for a in args[1:]]
+    ]
+
+    return ret
+
+def set_channel(*args, n_channels=3):
+    def _set_channel(img):
+        if img.ndim == 2:
+            img = np.expand_dims(img, axis=2)
+
+        c = img.shape[2]
+        if n_channels == 1 and c == 3:
+            img = np.expand_dims(sc.rgb2ycbcr(img)[:, :, 0], 2)
+        elif n_channels == 3 and c == 1:
+            img = np.concatenate([img] * n_channels, 2)
+
+        return img
+
+    return [_set_channel(a) for a in args]
+
+def np2Tensor(*args, rgb_range=255):
+    def _np2Tensor(img):
+        np_transpose = np.ascontiguousarray(img.transpose((2, 0, 1)))
+        tensor = torch.from_numpy(np_transpose).float()
+        tensor.mul_(rgb_range / 255)
+
+        return tensor
+
+    return [_np2Tensor(a) for a in args]
+
+def augment(*args, hflip=True, rot=True):
+    hflip = hflip and random.random() < 0.5
+    vflip = rot and random.random() < 0.5
+    rot90 = rot and random.random() < 0.5
+    # hflip=True
+    # vflip=True
+    # rot90=True
+
+    def _augment(img):
+        if hflip: img = img[:, ::-1]
+        if vflip: img = img[::-1, :]
+        if rot90: 
+            # print(img.shape)
+            img = img.transpose(1, 0)
+        
+        return img
+
+    return [_augment(a) for a in args]
+

CAR/code/data/demo.py

@@ -0,0 +1,39 @@
+import os
+
+from data import common
+
+import numpy as np
+import imageio
+
+import torch
+import torch.utils.data as data
+
+class Demo(data.Dataset):
+    def __init__(self, args, name='Demo', train=False, benchmark=False):
+        self.args = args
+        self.name = name
+        self.scale = args.scale
+        self.idx_scale = 0
+        self.train = False
+        self.benchmark = benchmark
+
+        self.filelist = []
+        for f in os.listdir(args.dir_demo):
+            if f.find('.png') >= 0 or f.find('.jp') >= 0:
+                self.filelist.append(os.path.join(args.dir_demo, f))
+        self.filelist.sort()
+
+    def __getitem__(self, idx):
+        filename = os.path.splitext(os.path.basename(self.filelist[idx]))[0]
+        lr = imageio.imread(self.filelist[idx])
+        lr, = common.set_channel(lr, n_channels=self.args.n_colors)
+        lr_t, = common.np2Tensor(lr, rgb_range=self.args.rgb_range)
+
+        return lr_t, -1, filename
+
+    def __len__(self):
+        return len(self.filelist)
+
+    def set_scale(self, idx_scale):
+        self.idx_scale = idx_scale
+

CAR/code/data/div2k.py

@@ -0,0 +1,32 @@
+import os
+from data import srdata
+
+class DIV2K(srdata.SRData):
+    def __init__(self, args, name='DIV2K', train=True, benchmark=False):
+        data_range = [r.split('-') for r in args.data_range.split('/')]
+        if train:
+            data_range = data_range[0]
+        else:
+            if args.test_only and len(data_range) == 1:
+                data_range = data_range[0]
+            else:
+                data_range = data_range[1]
+
+        self.begin, self.end = list(map(lambda x: int(x), data_range))
+        super(DIV2K, self).__init__(
+            args, name=name, train=train, benchmark=benchmark
+        )
+
+    def _scan(self):
+        names_hr, names_lr = super(DIV2K, self)._scan()
+        names_hr = names_hr[self.begin - 1:self.end]
+        names_lr = [n[self.begin - 1:self.end] for n in names_lr]
+
+        return names_hr, names_lr
+
+    def _set_filesystem(self, dir_data):
+        super(DIV2K, self)._set_filesystem(dir_data)
+        self.dir_hr = os.path.join(self.apath, 'DIV2K_train_HR')
+        self.dir_lr = os.path.join(self.apath, 'DIV2K_train_LR_bicubic')
+        if self.input_large: self.dir_lr += 'L'
+

CAR/code/data/div2kjpeg.py

@@ -0,0 +1,20 @@
+import os
+from data import srdata
+from data import div2k
+
+class DIV2KJPEG(div2k.DIV2K):
+    def __init__(self, args, name='', train=True, benchmark=False):
+        self.q_factor = int(name.replace('DIV2K-Q', ''))
+        super(DIV2KJPEG, self).__init__(
+            args, name=name, train=train, benchmark=benchmark
+        )
+
+    def _set_filesystem(self, dir_data):
+        self.apath = os.path.join(dir_data, 'DIV2K')
+        self.dir_hr = os.path.join(self.apath, 'DIV2K_train_HR')
+        self.dir_lr = os.path.join(
+            self.apath, 'DIV2K_Q{}'.format(self.q_factor)
+        )
+        if self.input_large: self.dir_lr += 'L'
+        self.ext = ('.png', '.jpg')
+

CAR/code/data/sr291.py

@@ -0,0 +1,6 @@
+from data import srdata
+
+class SR291(srdata.SRData):
+    def __init__(self, args, name='SR291', train=True, benchmark=False):
+        super(SR291, self).__init__(args, name=name)
+

CAR/code/data/srdata.py

@@ -0,0 +1,166 @@
+import os
+import glob
+import random
+import pickle
+
+from data import common
+
+import numpy as np
+import imageio
+import torch
+import torch.utils.data as data
+
+class SRData(data.Dataset):
+    def __init__(self, args, name='', train=True, benchmark=False):
+        self.args = args
+        self.name = name
+        self.train = train
+        self.split = 'train' if train else 'test'
+        self.do_eval = True
+        self.benchmark = benchmark
+        self.input_large = (args.model == 'VDSR')
+        self.scale = args.scale
+        self.idx_scale = 0
+        
+        self._set_filesystem(args.dir_data)
+        if args.ext.find('img') < 0:
+            path_bin = os.path.join(self.apath, 'bin')
+            os.makedirs(path_bin, exist_ok=True)
+
+        list_hr, list_lr = self._scan()
+        if args.ext.find('img') >= 0 or benchmark:
+            self.images_hr, self.images_lr = list_hr, list_lr
+        elif args.ext.find('sep') >= 0:
+            os.makedirs(
+                self.dir_hr.replace(self.apath, path_bin),
+                exist_ok=True
+            )
+            for s in self.scale:
+                os.makedirs(
+                    os.path.join(
+                        self.dir_lr.replace(self.apath, path_bin),
+                        'X{}'.format(s)
+                    ),
+                    exist_ok=True
+                )
+            
+            self.images_hr, self.images_lr = [], [[] for _ in self.scale]
+            for h in list_hr:
+                b = h.replace(self.apath, path_bin)
+                b = b.replace(self.ext[0], '.pt')
+                self.images_hr.append(b)
+                self._check_and_load(args.ext, h, b, verbose=True) 
+            for i, ll in enumerate(list_lr):
+                for l in ll:
+                    b = l.replace(self.apath, path_bin)
+                    b = b.replace(self.ext[1], '.pt')
+                    self.images_lr[i].append(b)
+                    self._check_and_load(args.ext, l, b, verbose=True) 
+        if train:
+            n_patches = args.batch_size * args.test_every
+            n_images = len(args.data_train) * len(self.images_hr)
+            if n_images == 0:
+                self.repeat = 0
+            else:
+                self.repeat = max(n_patches // n_images, 1)
+
+    # Below functions as used to prepare images
+    def _scan(self):
+        names_hr = sorted(
+            glob.glob(os.path.join(self.dir_hr, '*' + self.ext[0]))
+        )
+        names_lr = [[] for _ in self.scale]
+        for f in names_hr:
+            filename, _ = os.path.splitext(os.path.basename(f))
+            for si, s in enumerate(self.scale):
+                names_lr[si].append(os.path.join(
+                    self.dir_lr, 'X{}/{}{}'.format(
+                        s, filename, self.ext[1]
+                    )
+                ))
+
+        return names_hr, names_lr
+
+    def _set_filesystem(self, dir_data):
+        self.apath = os.path.join(dir_data, self.name)
+        self.dir_hr = os.path.join(self.apath, 'HR')
+        self.dir_lr = os.path.join(self.apath, 'LR_bicubic')
+        if self.input_large: self.dir_lr += 'L'
+        self.ext = ('.png', '.jpg')
+
+    def _check_and_load(self, ext, img, f, verbose=True):
+        if not os.path.isfile(f) or ext.find('reset') >= 0:
+            if verbose:
+                print('Making a binary: {}'.format(f))
+            with open(f, 'wb') as _f:
+                pickle.dump(imageio.imread(img), _f)
+
+    def __getitem__(self, idx):
+        lr, hr, filename = self._load_file(idx)
+        pair = self.get_patch(lr, hr)
+        pair = common.set_channel(*pair, n_channels=self.args.n_colors)
+        pair_t = common.np2Tensor(*pair, rgb_range=self.args.rgb_range)
+
+        return pair_t[0], pair_t[1], filename
+
+    def __len__(self):
+        if self.train:
+            return len(self.images_hr) * self.repeat
+        else:
+            return len(self.images_hr)
+
+    def _get_index(self, idx):
+        if self.train:
+            return idx % len(self.images_hr)
+        else:
+            return idx
+
+    def _load_file(self, idx):
+        idx = self._get_index(idx)
+        f_hr = self.images_hr[idx]
+        f_lr = self.images_lr[self.idx_scale][idx]
+
+        filename, _ = os.path.splitext(os.path.basename(f_hr))
+        if self.args.ext == 'img' or self.benchmark:
+            hr = imageio.imread(f_hr)
+            lr = imageio.imread(f_lr)
+        elif self.args.ext.find('sep') >= 0:
+            with open(f_hr, 'rb') as _f:
+                hr = pickle.load(_f)
+            with open(f_lr, 'rb') as _f:
+                lr = pickle.load(_f)
+
+        return lr, hr, filename
+
+    def get_patch(self, lr, hr):
+        scale = self.scale[self.idx_scale]
+        if self.train:
+            lr, hr = common.get_patch(
+                lr, hr,
+                patch_size=self.args.patch_size,
+                scale=scale,
+                multi=(len(self.scale) > 1),
+                input_large=self.input_large
+            )
+            if not self.args.no_augment: lr, hr = common.augment(lr, hr)
+        else:
+            if self.args.model=="RECURSIONNET":
+                lr, hr = common.get_patch(
+                    lr, hr,
+                    patch_size=self.args.patch_size,
+                    scale=scale,
+                    multi=(len(self.scale) > 1),
+                    input_large=self.input_large
+                )
+            else:
+                ih, iw = lr.shape[:2]
+                hr = hr[0:ih * scale, 0:iw * scale]
+
+        return lr, hr
+
+    def set_scale(self, idx_scale):
+        if not self.input_large:
+            self.idx_scale = idx_scale
+        else:
+            self.idx_scale = random.randint(0, len(self.scale) - 1)
+

CAR/code/data/video.py

@@ -0,0 +1,44 @@
+import os
+
+from data import common
+
+import cv2
+import numpy as np
+import imageio
+
+import torch
+import torch.utils.data as data
+
+class Video(data.Dataset):
+    def __init__(self, args, name='Video', train=False, benchmark=False):
+        self.args = args
+        self.name = name
+        self.scale = args.scale
+        self.idx_scale = 0
+        self.train = False
+        self.do_eval = False
+        self.benchmark = benchmark
+
+        self.filename, _ = os.path.splitext(os.path.basename(args.dir_demo))
+        self.vidcap = cv2.VideoCapture(args.dir_demo)
+        self.n_frames = 0
+        self.total_frames = int(self.vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    def __getitem__(self, idx):
+        success, lr = self.vidcap.read()
+        if success:
+            self.n_frames += 1
+            lr, = common.set_channel(lr, n_channels=self.args.n_colors)
+            lr_t, = common.np2Tensor(lr, rgb_range=self.args.rgb_range)
+
+            return lr_t, -1, '{}_{:0>5}'.format(self.filename, self.n_frames)
+        else:
+            vidcap.release()
+            return None
+
+    def __len__(self):
+        return self.total_frames
+
+    def set_scale(self, idx_scale):
+        self.idx_scale = idx_scale
+

CAR/code/dataloader.py

@@ -0,0 +1,158 @@
+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)
+

CAR/code/demo.sb

@@ -0,0 +1,6 @@
+#!/bin/bash
+# PANet Train
+#python main.py --n_GPUs 2 --batch_size 16 --lr 1e-4 --decay 200-400-600-800 ---save_models --model PANET --scale 10 --patch_size 48 --save PANET_Q10 --n_feats 64 --data_train DIV2K --chop
+
+# Test
+python main.py --model PANET --save_results --n_GPUs 1 --chop --data_test classic5+LIVE1 --scale 40 --n_resblocks 80 --n_feats 64 --pre_train ../Q40.pt --test_only

CAR/code/demo.sh

@@ -0,0 +1,56 @@
+# EDSR baseline model (x2) + JPEG augmentation
+python main.py --n_GPUs 2 --batch_size 16 --reset  --save_models --model EDSRL1T --scale 1 --patch_size 48 --save EDSR_R16F64P48_CSAL1TK20 --n_feats 64 --data_train DIV2K --chop
+#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 --model EDSRL1 --save_results --n_GPUs 2 --chop --data_test Kodak24+CBSD68+Urban100 --scale 1 --pre_train ../experiment/EDSR_R16F64P48_CSALK20/model/model_best.pt --test_only
+
+#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
+

CAR/code/lambda_networks/__init__.py

@@ -0,0 +1,4 @@
+from lambda_networks.lambda_networks import LambdaLayer
+from lambda_networks.lambda_networks import Recursion
+from lambda_networks.rlambda_networks import RLambdaLayer
+λLayer = LambdaLayer

CAR/code/lambda_networks/lambda_networks.py

@@ -0,0 +1,161 @@
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+from einops import rearrange
+
+# my layer normalization
+
+class LayerNorm(nn.Module):
+    def __init__(self, eps= 1e-5):
+        super(LayerNorm, self).__init__()
+        self.eps = eps
+    def forward(self, input):
+        shape=tuple(input.size()[1:])
+        return F.layer_norm(input, shape, eps=self.eps)
+    def extra_repr(self):
+        return f'eps={self.eps}'
+
+# helpers functions
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+# lambda layer
+
+class LambdaLayer(nn.Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        dim_k,
+        n = None,
+        r = None,
+        heads = 4,
+        dim_out = None,
+        dim_u = 1,
+        normalization="batch"):
+        super().__init__()
+        dim_out = default(dim_out, dim)
+        self.u = dim_u # intra-depth dimension
+        self.heads = heads
+
+        assert (dim_out % heads) == 0, 'values dimension must be divisible by number of heads for multi-head query'
+        dim_v = dim_out // heads
+
+        self.to_q = nn.Conv2d(dim, dim_k * heads, 1, bias = False)
+        self.to_k = nn.Conv2d(dim, dim_k * dim_u, 1, bias = False)
+        self.to_v = nn.Conv2d(dim, dim_v * dim_u, 1, bias = False)
+        print(f"using {normalization} in lambda layer")
+        if normalization=="none":
+            self.norm_q = nn.Identity(dim_k * heads)
+            self.norm_v = nn.Identity(dim_v * dim_u)
+        elif normalization=="instance":
+            self.norm_q = nn.InstanceNorm2d(dim_k * heads)
+            self.norm_v = nn.InstanceNorm2d(dim_v * dim_u)
+        elif normalization=="layer":
+            self.norm_q = LayerNorm()
+            self.norm_v = LayerNorm()
+        else:
+            self.norm_q = nn.BatchNorm2d(dim_k * heads)
+            self.norm_v = nn.BatchNorm2d(dim_v * dim_u)
+            print(f"using BN in lambda layer?")
+
+        self.local_contexts = exists(r)
+        if exists(r):
+            assert (r % 2) == 1, 'Receptive kernel size should be odd'
+            self.pos_conv = nn.Conv3d(dim_u, dim_k, (1, r, r), padding = (0, r // 2, r // 2))
+        else:
+            assert exists(n), 'You must specify the total sequence length (h x w)'
+            self.pos_emb = nn.Parameter(torch.randn(n, n, dim_k, dim_u))
+
+
+    def forward(self, x):
+        b, c, hh, ww, u, h = *x.shape, self.u, self.heads
+
+        q = self.to_q(x)
+        k = self.to_k(x)
+        v = self.to_v(x)
+
+        q = self.norm_q(q)
+        v = self.norm_v(v)
+
+        q = rearrange(q, 'b (h k) hh ww -> b h k (hh ww)', h = h)
+        k = rearrange(k, 'b (u k) hh ww -> b u k (hh ww)', u = u)
+        v = rearrange(v, 'b (u v) hh ww -> b u v (hh ww)', u = u)
+
+        k = k.softmax(dim=-1)
+
+        λc = einsum('b u k m, b u v m -> b k v', k, v)
+        Yc = einsum('b h k n, b k v -> b h v n', q, λc)
+
+        if self.local_contexts:
+            v = rearrange(v, 'b u v (hh ww) -> b u v hh ww', hh = hh, ww = ww)
+            λp = self.pos_conv(v)
+            Yp = einsum('b h k n, b k v n -> b h v n', q, λp.flatten(3))
+        else:
+            λp = einsum('n m k u, b u v m -> b n k v', self.pos_emb, v)
+            Yp = einsum('b h k n, b n k v -> b h v n', q, λp)
+
+        Y = Yc + Yp
+        out = rearrange(Y, 'b h v (hh ww) -> b (h v) hh ww', hh = hh, ww = ww)
+        return out
+
+
+# i'm not sure whether this will work or not
+class Recursion(nn.Module):
+    def __init__(self, N: int, hidden_dim:int=64):
+        super(Recursion,self).__init__()
+        self.N = N
+        self.lambdaNxN_identity = LambdaLayer(dim=hidden_dim, dim_out=hidden_dim, n=N * N, dim_k=16, heads=2, dim_u=1)
+        # merge upstream information here
+        self.lambdaNxN_merge = LambdaLayer(dim=2*hidden_dim, dim_out=hidden_dim, n=N * N, dim_k=16, heads=2, dim_u=1)
+        self.downscale_conv = nn.Conv2d(hidden_dim, hidden_dim, kernel_size=N, stride=N)
+        self.upscale_conv = nn.Conv2d(hidden_dim, hidden_dim * N * N, kernel_size=3,padding=1)
+        self.pixel_shuffle = nn.PixelShuffle(N)
+    
+    def forward(self, x: torch.Tensor):
+        N = self.N
+
+        def to_patch(blocks:torch.Tensor)->torch.Tensor:
+            shape = blocks.shape
+            blocks_patch = F.unfold(blocks, kernel_size=N, stride=N)
+            blocks_patch = blocks_patch.view(shape[0], shape[1], N, N, -1)
+            num_patch = blocks_patch.shape[-1]
+            blocks_patch = blocks_patch.permute(0, 4, 1, 2, 3).reshape(-1, shape[1], N, N).contiguous()
+            return blocks_patch, num_patch
+        
+        def combine_patch(processed_patch,shape,num_patch):
+            processed_patch = processed_patch.reshape(shape[0], num_patch, shape[1], N, N)
+            processed_patch=processed_patch.permute(0, 2, 3, 4, 1).reshape(shape[0],shape[1] * N * N,num_patch).contiguous()
+            processed=F.fold(processed_patch,output_size=(shape[-2],shape[-1]),kernel_size=N,stride=N)
+            return processed
+
+        def process(blocks:torch.Tensor)->torch.Tensor:
+            shape = blocks.shape
+            if blocks.shape[-1] == N:
+                processed = self.lambdaNxN_identity(blocks)
+                return processed
+            # to NxN patchs
+            blocks_patch,num_patch=to_patch(blocks)
+            # pass through identity
+            processed_patch = self.lambdaNxN_identity(blocks_patch)
+            # back to HxW
+            processed=combine_patch(processed_patch,shape,num_patch)
+            # get feedback
+            feedback = process(self.downscale_conv(processed))
+            # upscale feedback
+            upscale_feedback = self.upscale_conv(feedback)
+            upscale_feedback=self.pixel_shuffle(upscale_feedback)
+            # combine results
+            combined = torch.cat([processed, upscale_feedback], dim=1)
+            combined_shape=combined.shape
+            combined_patch,num_patch=to_patch(combined)
+            combined_patch_reduced = self.lambdaNxN_merge(combined_patch)
+            ret_shape=(combined_shape[0],combined_shape[1]//2,combined_shape[2],combined_shape[3])
+            ret=combine_patch(combined_patch_reduced,ret_shape,num_patch)
+            return ret
+
+        return process(x)

CAR/code/lambda_networks/rlambda_networks.py

@@ -0,0 +1,93 @@
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+from einops import rearrange
+
+
+# helpers functions
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+# lambda layer
+
+class RLambdaLayer(nn.Module):
+    def __init__(
+            self,
+            dim,
+            *,
+            dim_k,
+            n=None,
+            r=None,
+            heads=4,
+            dim_out=None,
+            dim_u=1,
+            recurrence=None
+    ):
+        super().__init__()
+        dim_out = default(dim_out, dim)
+        self.u = dim_u  # intra-depth dimension
+        self.heads = heads
+
+        assert (dim_out % heads) == 0, 'values dimension must be divisible by number of heads for multi-head query'
+        dim_v = dim_out // heads
+
+        self.to_q = nn.Conv2d(dim, dim_k * heads, 1, bias=False)
+        self.to_k = nn.Conv2d(dim, dim_k * dim_u, 1, bias=False)
+        self.to_v = nn.Conv2d(dim, dim_v * dim_u, 1, bias=False)
+
+        self.norm_q = nn.BatchNorm2d(dim_k * heads)
+        self.norm_v = nn.BatchNorm2d(dim_v * dim_u)
+
+        self.local_contexts = exists(r)
+        self.recurrence = recurrence
+        if exists(r):
+            assert (r % 2) == 1, 'Receptive kernel size should be odd'
+            self.pos_conv = nn.Conv3d(dim_u, dim_k, (1, r, r), padding=(0, r // 2, r // 2))
+        else:
+            assert exists(n), 'You must specify the total sequence length (h x w)'
+            self.pos_emb = nn.Parameter(torch.randn(n, n, dim_k, dim_u))
+
+    def apply_lambda(self, lambda_c, lambda_p, x):
+        b, c, hh, ww, u, h = *x.shape, self.u, self.heads
+        q = self.to_q(x)
+        q = self.norm_q(q)
+        q = rearrange(q, 'b (h k) hh ww -> b h k (hh ww)', h=h)
+        Yc = einsum('b h k n, b k v -> b h v n', q, lambda_c)
+        if self.local_contexts:
+            Yp = einsum('b h k n, b k v n -> b h v n', q, lambda_p.flatten(3))
+        else:
+            Yp = einsum('b h k n, b n k v -> b h v n', q, lambda_p)
+        Y = Yc + Yp
+        out = rearrange(Y, 'b h v (hh ww) -> b (h v) hh ww', hh=hh, ww=ww)
+        return out
+
+    def forward(self, x):
+        b, c, hh, ww, u, h = *x.shape, self.u, self.heads
+
+        k = self.to_k(x)
+        v = self.to_v(x)
+
+        v = self.norm_v(v)
+
+        k = rearrange(k, 'b (u k) hh ww -> b u k (hh ww)', u=u)
+        v = rearrange(v, 'b (u v) hh ww -> b u v (hh ww)', u=u)
+
+        k = k.softmax(dim=-1)
+
+        λc = einsum('b u k m, b u v m -> b k v', k, v)
+
+        if self.local_contexts:
+            v = rearrange(v, 'b u v (hh ww) -> b u v hh ww', hh=hh, ww=ww)
+            λp = self.pos_conv(v)
+        else:
+            λp = einsum('n m k u, b u v m -> b n k v', self.pos_emb, v)
+        out = x
+        for i in range(self.recurrence):
+            out = self.apply_lambda(λc, λp, out)
+        return out

CAR/code/loss/__init__.py

@@ -0,0 +1,173 @@
+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
+
+def sequence_loss(sr, hr, loss_func, gamma=0.8, max_val=None):
+    """ Loss function defined over sequence of flow predictions """
+
+    n_recurrence = len(sr)    
+    total_loss = 0.0
+    buffer=[0.0]*n_recurrence
+    # exlude invalid pixels and extremely large diplacements
+    for i in range(n_recurrence):
+        i_weight = gamma**(n_recurrence - i - 1)
+        i_loss = loss_func(sr[i],hr)
+        buffer[i]=i_loss.item()
+        # total_loss += i_weight * (valid[:, None] * i_loss).mean()
+        total_loss += i_weight * (i_loss)
+    return total_loss,buffer
+
+class Loss(nn.modules.loss._Loss):
+    def __init__(self, args, ckp):
+        super(Loss, self).__init__()
+        print('Preparing loss function:')
+        self.buffer=[0.0]*args.recurrence
+        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:
+                if isinstance(sr,list):
+                    # weights=[0.32,0.08,0.02,0.01,0.005]
+                    # weights=weights[::-1]
+                    # weights=[0.01,0.02,0.08,0.32]
+                    # self.buffer=[]
+                    effective_loss,buffer_lst=sequence_loss(sr,hr,l['function'])
+                    # for k in range(len(sr)):
+                        # loss=l['function'](sr[k], hr)
+                        # self.buffer.append(loss.item())
+                        # effective_loss=loss*weights[k]*l['weight']
+                    losses.append(effective_loss)
+                    self.buffer=buffer_lst
+                    self.log[-1, i] += effective_loss.item()
+                else:
+                    loss = l['function'](sr, hr)
+                    effective_loss = l['weight'] * loss
+                    losses.append(effective_loss)
+                    self.buffer[0]=effective_loss.item()
+                    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()
+

CAR/code/loss/adversarial.py

@@ -0,0 +1,112 @@
+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

CAR/code/loss/discriminator.py

@@ -0,0 +1,55 @@
+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
+

CAR/code/loss/vgg.py

@@ -0,0 +1,36 @@
+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

CAR/code/main.py

@@ -0,0 +1,35 @@
+import torch
+
+import utility
+import data
+import model
+import loss
+from option import args
+from trainer import Trainer
+
+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)
+        print('total params: %.2fM' % (sum(p.numel() for p in model.parameters())/1000000.0))
+        t = VideoTester(args, model, checkpoint)
+        t.test()
+    else:
+        if checkpoint.ok:
+            loader = data.Data(args)
+            _model = model.Model(args, checkpoint)
+            print('total params:%.5fM' % (sum(p.numel() for p in _model.parameters())/1000000.0))
+            _loss = loss.Loss(args, checkpoint) if not args.test_only else None
+            t = Trainer(args, loader, _model, _loss, checkpoint)
+            while not t.terminate():
+                t.train()
+                t.test()
+
+            checkpoint.done()
+
+if __name__ == '__main__':
+    main()

CAR/code/model/LICENSE

@@ -0,0 +1,21 @@
+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.

CAR/code/model/__init__.py

@@ -0,0 +1,190 @@
+import os
+from importlib import import_module
+
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+
+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.self_ensemble = args.self_ensemble
+        self.chop = args.chop
+        self.precision = args.precision
+        self.cpu = args.cpu
+        self.device = torch.device('cpu' if args.cpu else 'cuda')
+        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()
+
+        if not args.cpu and args.n_GPUs > 1:
+            self.model = nn.DataParallel(self.model, range(args.n_GPUs))
+
+        self.load(
+            ckp.dir,
+            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
+        target = self.get_model()
+        if hasattr(target, 'set_scale'):
+            target.set_scale(idx_scale)
+
+        if self.self_ensemble and not self.training:
+            if self.chop:
+                forward_function = self.forward_chop
+            else:
+                forward_function = self.model.forward
+
+            return self.forward_x8(x, forward_function)
+        elif self.chop and not self.training:
+            return self.forward_chop(x)
+        else:
+            return self.model(x)
+
+    def get_model(self):
+        if self.n_GPUs == 1:
+            return self.model
+        else:
+            return self.model.module
+
+    def state_dict(self, **kwargs):
+        target = self.get_model()
+        return target.state_dict(**kwargs)
+
+    def save(self, apath, epoch, is_best=False):
+        target = self.get_model()
+        torch.save(
+            target.state_dict(), 
+            os.path.join(apath, 'model_latest.pt')
+        )
+        if is_best:
+            torch.save(
+                target.state_dict(),
+                os.path.join(apath, 'model_best.pt')
+            )
+        
+        if self.save_models:
+            torch.save(
+                target.state_dict(),
+                os.path.join(apath, 'model_{}.pt'.format(epoch))
+            )
+
+    def load(self, apath, pre_train='.', resume=-1, cpu=False):
+        if cpu:
+            kwargs = {'map_location': lambda storage, loc: storage}
+        else:
+            kwargs = {}
+
+        if resume == -1:
+            self.get_model().load_state_dict(
+                torch.load(
+                    os.path.join(apath,'model', 'model_latest.pt'),
+                    **kwargs
+                ),
+                strict=False
+            )
+        elif resume == 0:
+            if pre_train != '.':
+                print('Loading model from {}'.format(pre_train))
+                self.get_model().load_state_dict(
+                    torch.load(pre_train, **kwargs),
+                    strict=False
+                )
+        else:
+            self.get_model().load_state_dict(
+                torch.load(
+                    os.path.join(apath, 'model', 'model_{}.pt'.format(resume)),
+                    **kwargs
+                ),
+                strict=False
+            )
+
+    def forward_chop(self, x, shave=10, min_size=6400):
+        scale = self.scale[self.idx_scale]
+        scale = 1
+        n_GPUs = min(self.n_GPUs, 4)
+        b, c, h, w = x.size()
+        h_half, w_half = h // 2, w // 2
+        h_size, w_size = h_half + shave, w_half + shave
+        lr_list = [
+            x[:, :, 0:h_size, 0:w_size],
+            x[:, :, 0:h_size, (w - w_size):w],
+            x[:, :, (h - h_size):h, 0:w_size],
+            x[:, :, (h - h_size):h, (w - w_size):w]]
+
+        if w_size * h_size < min_size:
+            sr_list = []
+            for i in range(0, 4, n_GPUs):
+                lr_batch = torch.cat(lr_list[i:(i + n_GPUs)], dim=0)
+                sr_batch = self.model(lr_batch)
+                sr_list.extend(sr_batch.chunk(n_GPUs, dim=0))
+        else:
+            sr_list = [
+                self.forward_chop(patch, shave=shave, min_size=min_size) \
+                for patch in lr_list
+            ]
+
+        h, w = scale * h, scale * w
+        h_half, w_half = scale * h_half, scale * w_half
+        h_size, w_size = scale * h_size, scale * w_size
+        shave *= scale
+
+        output = x.new(b, c, h, w)
+        output[:, :, 0:h_half, 0:w_half] \
+            = sr_list[0][:, :, 0:h_half, 0:w_half]
+        output[:, :, 0:h_half, w_half:w] \
+            = sr_list[1][:, :, 0:h_half, (w_size - w + w_half):w_size]
+        output[:, :, h_half:h, 0:w_half] \
+            = sr_list[2][:, :, (h_size - h + h_half):h_size, 0:w_half]
+        output[:, :, h_half:h, w_half:w] \
+            = sr_list[3][:, :, (h_size - h + h_half):h_size, (w_size - w + w_half):w_size]
+
+        return output
+
+    def forward_x8(self, x, forward_function):
+        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
+
+        lr_list = [x]
+        for tf in 'v', 'h', 't':
+            lr_list.extend([_transform(t, tf) for t in lr_list])
+
+        sr_list = [forward_function(aug) for aug in lr_list]
+        for i in range(len(sr_list)):
+            if i > 3:
+                sr_list[i] = _transform(sr_list[i], 't')
+            if i % 4 > 1:
+                sr_list[i] = _transform(sr_list[i], 'h')
+            if (i % 4) % 2 == 1:
+                sr_list[i] = _transform(sr_list[i], 'v')
+
+        output_cat = torch.cat(sr_list, dim=0)
+        output = output_cat.mean(dim=0, keepdim=True)
+
+        return output
+