Upload folder using huggingface_hub

SR/README.md

@@ -0,0 +1,78 @@
+# 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.
+
+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: 
+    python main.py --n_GPUs 4 --rgb_range 1 --reset --save_models --lr 1e-4 --decay 200-400-600-800 --chop --save_results --n_resblocks 32 --n_feats 256 --res_scale 0.1 --batch_size 16 --model PAEDSR --scale 2 --patch_size 96 --save EDSR_PA_x2 --data_train DIV2K
+    ```
+## 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 to reproduce the results in the paper.
+    # Example Usage: 
+    python main.py --model PAEDSR --data_test Set5+Set14+B100+Urban100 --save_results --rgb_range 1 --data_range 801-900 --scale 2 --n_feats 256 --n_resblocks 32 --res_scale 0.1 --pre_train ../model_x2.pt --test_only --chop 
+
+    ```
+
+### The whole test pipeline
+1. Prepare benchmark datasets [from SNU_CVLab](https://cv.snu.ac.kr/research/EDSR/benchmark.tar) 
+
+2. Conduct image SR. 
+
+    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.

SR/code/.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python.pythonPath": "D:\\Environments\\Anaconda3\\envs\\torch1.6\\python.exe"
+}

SR/code/bootstrap.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+
+wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
+
+bash ~/miniconda.sh -b -p $HOME/miniconda

SR/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 ['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,
+                )
+            )

SR/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 = ('', '.png')
+

SR/code/data/common.py

@@ -0,0 +1,72 @@
+import random
+
+import numpy as np
+import skimage.color as sc
+
+import torch
+
+def get_patch(*args, patch_size=96, scale=2, multi=False, input_large=False):
+    ih, iw = args[0].shape[:2]
+
+    if not input_large:
+        p = scale if multi else 1
+        tp = p * patch_size
+        ip = tp // scale
+    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 = scale * ix, scale * 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
+
+    def _augment(img):
+        if hflip: img = img[:, ::-1, :]
+        if vflip: img = img[::-1, :, :]
+        if rot90: img = img.transpose(1, 0, 2)
+        
+        return img
+
+    return [_augment(a) for a in args]
+

SR/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
+

SR/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'
+

SR/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')
+

SR/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)
+

SR/code/data/srdata.py

@@ -0,0 +1,157 @@
+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{}/{}x{}{}'.format(
+                        s, filename, s, 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', '.png')
+
+    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:
+            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)
+

SR/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
+

SR/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)
+

SR/code/demo.sb

@@ -0,0 +1,5 @@
+#!/bin/bash
+#Train x2
+python main.py --n_GPUs 4 --rgb_range 1 --reset --save_models --lr 1e-4 --decay 200-400-600-800 --chop --save_results --n_resblocks 32 --n_feats 256 --res_scale 0.1 --batch_size 16 --model PAEDSR --scale 2 --patch_size 96 --save EDSR_PA_x2 --data_train DIV2K
+#Test
+python main.py --model PAEDSR --data_test Set5+Set14+B100+Urban100 --save_results --rgb_range 1 --data_range 801-900 --scale 2 --n_feats 256 --n_resblocks 32 --res_scale 0.1 --pre_train ../model_x2.pt --test_only --chop 

SR/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

SR/code/lambda_networks/lambda_networks.py

@@ -0,0 +1,137 @@
+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 LambdaLayer(nn.Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        dim_k,
+        n = None,
+        r = None,
+        heads = 4,
+        dim_out = None,
+        dim_u = 1):
+        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)
+        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)

SR/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

SR/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()
+