Upload folder using huggingface_hub

deblur/experiment/.gitignore

@@ -0,0 +1 @@
+!.gitignore

deblur/src/data/__init__.py

@@ -0,0 +1,79 @@
+"""Generic dataset loader"""
+
+from importlib import import_module
+
+from torch.utils.data import DataLoader
+from torch.utils.data import SequentialSampler, RandomSampler
+from torch.utils.data.distributed import DistributedSampler
+from .sampler import DistributedEvalSampler
+
+class Data():
+    def __init__(self, args):
+
+        self.modes = ['train', 'val', 'test', 'demo']
+
+        self.action = {
+            'train': args.do_train,
+            'val':  args.do_validate,
+            'test': args.do_test,
+            'demo': args.demo
+        }
+
+        self.dataset_name = {
+            'train': args.data_train,
+            'val': args.data_val,
+            'test': args.data_test,
+            'demo': 'Demo'
+        }
+
+        self.args = args
+
+        def _get_data_loader(mode='train'):
+            dataset_name = self.dataset_name[mode]
+            dataset = import_module('data.' + dataset_name.lower())
+            dataset = getattr(dataset, dataset_name)(args, mode)
+
+            if mode == 'train':
+                if args.distributed:
+                    batch_size = int(args.batch_size / args.n_GPUs)   # batch size per GPU (single-node training)
+                    sampler = DistributedSampler(dataset, shuffle=True, num_replicas=args.world_size, rank=args.rank)
+                    num_workers = int((args.num_workers + args.n_GPUs - 1) / args.n_GPUs)    # num_workers per GPU (single-node training)
+                else:
+                    batch_size = args.batch_size
+                    sampler = RandomSampler(dataset, replacement=False)
+                    num_workers = args.num_workers
+                drop_last = True
+
+            elif mode in ('val', 'test', 'demo'):
+                if args.distributed:
+                    batch_size = 1  # 1 image per GPU
+                    sampler = DistributedEvalSampler(dataset, shuffle=False, num_replicas=args.world_size, rank=args.rank)
+                    num_workers = int((args.num_workers + args.n_GPUs - 1) / args.n_GPUs)    # num_workers per GPU (single-node training)
+                else:
+                    batch_size = args.n_GPUs    # 1 image per GPU
+                    sampler = SequentialSampler(dataset)
+                    num_workers = args.num_workers
+                drop_last = False
+
+            loader = DataLoader(
+                dataset=dataset,
+                batch_size=batch_size,
+                shuffle=False,
+                sampler=sampler,
+                num_workers=num_workers,
+                pin_memory=True,
+                drop_last=drop_last,
+            )
+
+            return loader
+
+        self.loaders = {}
+        for mode in self.modes:
+            if self.action[mode]:
+                self.loaders[mode] = _get_data_loader(mode)
+                print('===> Loading {} dataset: {}'.format(mode, self.dataset_name[mode]))
+            else:
+                self.loaders[mode] = None
+
+    def get_loader(self):
+        return self.loaders

deblur/src/data/common.py

@@ -0,0 +1,165 @@
+import random
+import numpy as np
+from skimage.color import rgb2hsv, hsv2rgb
+# from skimage.transform import pyramid_gaussian
+
+import torch
+
+def _apply(func, x):
+
+    if isinstance(x, (list, tuple)):
+        return [_apply(func, x_i) for x_i in x]
+    elif isinstance(x, dict):
+        y = {}
+        for key, value in x.items():
+            y[key] = _apply(func, value)
+        return y
+    else:
+        return func(x)
+
+def crop(*args, ps=256):    # patch_size
+    # args = [input, target]
+    def _get_shape(*args):
+        if isinstance(args[0], (list, tuple)):
+            return _get_shape(args[0][0])
+        elif isinstance(args[0], dict):
+            return _get_shape(list(args[0].values())[0])
+        else:
+            return args[0].shape
+
+    h, w, _ = _get_shape(args)
+
+    py = random.randrange(0, h-ps+1)
+    px = random.randrange(0, w-ps+1)
+
+    def _crop(img):
+        if img.ndim == 2:
+            return img[py:py+ps, px:px+ps, np.newaxis]
+        else:
+            return img[py:py+ps, px:px+ps, :]
+
+    return _apply(_crop, args)
+
+def add_noise(*args, sigma_sigma=2, rgb_range=255):
+
+    if len(args) == 1:  # usually there is only a single input
+        args = args[0]
+
+    sigma = np.random.normal() * sigma_sigma * rgb_range/255
+
+    def _add_noise(img):
+        noise = np.random.randn(*img.shape).astype(np.float32) * sigma
+        return (img + noise).clip(0, rgb_range)
+
+    return _apply(_add_noise, args)
+
+def augment(*args, hflip=True, rot=True, shuffle=True, change_saturation=True, rgb_range=255):
+    """augmentation consistent to input and target"""
+
+    choices = (False, True)
+
+    hflip = hflip and random.choice(choices)
+    vflip = rot and random.choice(choices)
+    rot90 = rot and random.choice(choices)
+    # shuffle = shuffle
+
+    if shuffle:
+        rgb_order = list(range(3))
+        random.shuffle(rgb_order)
+        if rgb_order == list(range(3)):
+            shuffle = False
+
+    if change_saturation:
+        amp_factor = np.random.uniform(0.5, 1.5)
+
+    def _augment(img):
+        if hflip: img = img[:, ::-1, :]
+        if vflip: img = img[::-1, :, :]
+        if rot90: img = img.transpose(1, 0, 2)
+        if shuffle and img.ndim > 2:
+            if img.shape[-1] == 3:  # RGB image only
+                img = img[..., rgb_order]
+
+        if change_saturation:
+            hsv_img = rgb2hsv(img)
+            hsv_img[..., 1] *= amp_factor
+
+            img = hsv2rgb(hsv_img).clip(0, 1) * rgb_range
+
+        return img.astype(np.float32)
+
+    return _apply(_augment, args)
+
+def pad(img, divisor=4, pad_width=None, negative=False):
+
+    def _pad_numpy(img, divisor=4, pad_width=None, negative=False):
+        if pad_width is None:
+            (h, w, _) = img.shape
+            pad_h = -h % divisor
+            pad_w = -w % divisor
+            pad_width = ((0, pad_h), (0, pad_w), (0, 0))
+
+        img = np.pad(img, pad_width, mode='edge')
+
+        return img, pad_width
+
+    def _pad_tensor(img, divisor=4, pad_width=None, negative=False):
+
+        n, c, h, w = img.shape
+        if pad_width is None:
+            pad_h = -h % divisor
+            pad_w = -w % divisor
+            pad_width = (0, pad_w, 0, pad_h)
+        else:
+            try:
+                pad_h = pad_width[0][1]
+                pad_w = pad_width[1][1]
+                if isinstance(pad_h, torch.Tensor):
+                    pad_h = pad_h.item()
+                if isinstance(pad_w, torch.Tensor):
+                    pad_w = pad_w.item()
+
+                pad_width = (0, pad_w, 0, pad_h)
+            except:
+                pass
+
+            if negative:
+                pad_width = [-val for val in pad_width]
+
+        img = torch.nn.functional.pad(img, pad_width, 'reflect')
+
+        return img, pad_width
+
+    if isinstance(img, np.ndarray):
+        return _pad_numpy(img, divisor, pad_width, negative)
+    else:   # torch.Tensor
+        return _pad_tensor(img, divisor, pad_width, negative)
+
+def generate_pyramid(*args, n_scales):
+
+    def _generate_pyramid(img):
+        if img.dtype != np.float32:
+            img = img.astype(np.float32)
+        # pyramid = list(pyramid_gaussian(img, n_scales-1, multichannel=True))
+        # bypass pyramid, deliver the image as is
+        pyramid = [img]
+
+        return pyramid
+
+    return _apply(_generate_pyramid, args)
+
+def np2tensor(*args, rgb_range=255):
+    def _np2tensor(x):
+        np_transpose = np.ascontiguousarray(x.transpose(2, 0, 1))
+        tensor = torch.from_numpy(np_transpose).float()
+        tensor.mul_(rgb_range / 255)
+        return tensor
+
+    return _apply(_np2tensor, args)
+
+def to(*args, device=None, dtype=torch.float):
+
+    def _to(x):
+        return x.to(device=device, dtype=dtype, non_blocking=True, copy=False)
+
+    return _apply(_to, args)

deblur/src/data/dataset.py

@@ -0,0 +1,154 @@
+import os
+import random
+import imageio
+import numpy as np
+import torch.utils.data as data
+
+from data import common
+
+from utils import interact
+
+class Dataset(data.Dataset):
+    """Basic dataloader class
+    """
+    def __init__(self, args, mode='train'):
+        super(Dataset, self).__init__()
+        self.args = args
+        self.mode = mode
+
+        self.modes = ()
+        self.set_modes()
+        self._check_mode()
+
+        self.set_keys()
+
+        if self.mode == 'train':
+            dataset = args.data_train
+        elif self.mode == 'val':
+            dataset = args.data_val
+        elif self.mode == 'test':
+            dataset = args.data_test
+        elif self.mode == 'demo':
+            pass
+        else:
+            raise NotImplementedError('not implemented for this mode: {}!'.format(self.mode))
+
+        if self.mode == 'demo':
+            self.subset_root = args.demo_input_dir
+        else:
+            self.subset_root = os.path.join(args.data_root, dataset, self.mode)
+
+        self.blur_list = []
+        self.sharp_list = []
+
+        self._scan()
+
+    def set_modes(self):
+        self.modes = ('train', 'val', 'test', 'demo')
+
+    def _check_mode(self):
+        """Should be called in the child class __init__() after super
+        """
+        if self.mode not in self.modes:
+            raise NotImplementedError('mode error: not for {}'.format(self.mode))
+
+        return
+
+    def set_keys(self):
+        self.blur_key = 'blur'      # to be overwritten by child class
+        self.sharp_key = 'sharp'    # to be overwritten by child class
+
+        self.non_blur_keys = []
+        self.non_sharp_keys = []
+
+        return
+
+    def _scan(self, root=None):
+        """Should be called in the child class __init__() after super
+        """
+        if root is None:
+            root = self.subset_root
+
+        if self.blur_key in self.non_blur_keys:
+            self.non_blur_keys.remove(self.blur_key)
+        if self.sharp_key in self.non_sharp_keys:
+            self.non_sharp_keys.remove(self.sharp_key)
+
+        def _key_check(path, true_key, false_keys):
+            path = os.path.join(path, '')
+            if path.find(true_key) >= 0:
+                for false_key in false_keys:
+                    if path.find(false_key) >= 0:
+                        return False
+
+                return True
+            else:
+                return False
+
+        def _get_list_by_key(root, true_key, false_keys):
+            data_list = []
+            for sub, dirs, files in os.walk(root):
+                if not dirs:
+                    file_list = [os.path.join(sub, f) for f in files]
+                    if _key_check(sub, true_key, false_keys):
+                        data_list += file_list
+
+            data_list.sort()
+
+            return data_list
+
+        def _rectify_keys():
+            self.blur_key = os.path.join(self.blur_key, '')
+            self.non_blur_keys = [os.path.join(non_blur_key, '') for non_blur_key in self.non_blur_keys]
+            self.sharp_key = os.path.join(self.sharp_key, '')
+            self.non_sharp_keys = [os.path.join(non_sharp_key, '') for non_sharp_key in self.non_sharp_keys]
+
+        _rectify_keys()
+
+        self.blur_list = _get_list_by_key(root, self.blur_key, self.non_blur_keys)
+        self.sharp_list = _get_list_by_key(root, self.sharp_key, self.non_sharp_keys)
+
+        if len(self.sharp_list) > 0:
+            assert(len(self.blur_list) == len(self.sharp_list))
+
+        return
+
+    def __getitem__(self, idx):
+
+        blur = imageio.imread(self.blur_list[idx], pilmode='RGB')
+        if len(self.sharp_list) > 0:
+            sharp = imageio.imread(self.sharp_list[idx], pilmode='RGB')
+            imgs = [blur, sharp]
+        else:
+            imgs = [blur]
+
+        pad_width = 0   # dummy value
+        if self.mode == 'train':
+            imgs = common.crop(*imgs, ps=self.args.patch_size)
+            if self.args.augment:
+                imgs = common.augment(*imgs, hflip=True, rot=True, shuffle=True, change_saturation=True, rgb_range=self.args.rgb_range)
+                imgs[0] = common.add_noise(imgs[0], sigma_sigma=2, rgb_range=self.args.rgb_range)
+        elif self.mode == 'demo':
+            imgs[0], pad_width = common.pad(imgs[0], divisor=2**(self.args.n_scales-1))   # pad in case of non-divisible size
+        else:
+            pass    # deliver test image as is.
+
+        if self.args.gaussian_pyramid:
+            imgs = common.generate_pyramid(*imgs, n_scales=self.args.n_scales)
+
+        imgs = common.np2tensor(*imgs, rgb_range=self.args.rgb_range)
+        relpath = os.path.relpath(self.blur_list[idx], self.subset_root)
+
+        blur = imgs[0]
+        sharp = imgs[1] if len(imgs) > 1 else False
+
+        return blur, sharp, pad_width, idx, relpath
+
+    def __len__(self):
+        return len(self.blur_list)
+        # return 32
+
+
+
+
+

deblur/src/data/demo.py

@@ -0,0 +1,22 @@
+from data.dataset import Dataset
+
+from utils import interact
+
+class Demo(Dataset):
+    """Demo train, test subset class
+    """
+    def __init__(self, args, mode='demo'):
+        super(Demo, self).__init__(args, mode)
+
+    def set_modes(self):
+        self.modes = ('demo')
+
+    def set_keys(self):
+        super(Demo, self).set_keys()
+        self.blur_key = ''              # all the files
+        self.non_sharp_keys = ['']      # no files
+
+    def __getitem__(self, idx):
+        blur, sharp, pad_width, idx, relpath = super(Demo, self).__getitem__(idx)
+
+        return blur, sharp, pad_width, idx, relpath

deblur/src/data/gopro_large.py

@@ -0,0 +1,23 @@
+from data.dataset import Dataset
+
+from utils import interact
+
+class GOPRO_Large(Dataset):
+    """GOPRO_Large train, test subset class
+    """
+    def __init__(self, args, mode='train'):
+        super(GOPRO_Large, self).__init__(args, mode)
+
+    def set_modes(self):
+        self.modes = ('train', 'test')
+
+    def set_keys(self):
+        super(GOPRO_Large, self).set_keys()
+        self.blur_key = 'blur_gamma'
+        # self.sharp_key = 'sharp'
+
+    def __getitem__(self, idx):
+        blur, sharp, pad_width, idx, relpath = super(GOPRO_Large, self).__getitem__(idx)
+        relpath = relpath.replace('{}/'.format(self.blur_key), '')
+
+        return blur, sharp, pad_width, idx, relpath

deblur/src/data/reds.py

@@ -0,0 +1,28 @@
+from data.dataset import Dataset
+
+from utils import interact
+
+class REDS(Dataset):
+    """REDS train, val, test subset class
+    """
+    def __init__(self, args, mode='train'):
+        super(REDS, self).__init__(args, mode)
+
+    def set_modes(self):
+        self.modes = ('train', 'val', 'test')
+
+    def set_keys(self):
+        super(REDS, self).set_keys()
+        # self.blur_key = 'blur'
+        # self.sharp_key = 'sharp'
+
+        self.non_blur_keys = ['blur', 'blur_comp', 'blur_bicubic']
+        self.non_blur_keys.remove(self.blur_key)
+        self.non_sharp_keys = ['sharp_bicubic', 'sharp']
+        self.non_sharp_keys.remove(self.sharp_key)
+
+    def __getitem__(self, idx):
+        blur, sharp, pad_width, idx, relpath = super(REDS, self).__getitem__(idx)
+        relpath = relpath.replace('{}/{}/'.format(self.mode, self.blur_key), '')
+
+        return blur, sharp, pad_width, idx, relpath

deblur/src/data/sampler.py

@@ -0,0 +1,115 @@
+import math
+import torch
+from torch.utils.data import Sampler
+import torch.distributed as dist
+
+
+class DistributedEvalSampler(Sampler):
+    r"""
+    DistributedEvalSampler is different from DistributedSampler.
+    It does NOT add extra samples to make it evenly divisible.
+    DistributedEvalSampler should NOT be used for training. The distributed processes could hang forever.
+    See this issue for details: https://github.com/pytorch/pytorch/issues/22584
+    shuffle is disabled by default
+
+    DistributedEvalSampler is for evaluation purpose where synchronization does not happen every epoch.
+    Synchronization should be done outside the dataloader loop.
+
+    Sampler that restricts data loading to a subset of the dataset.
+
+    It is especially useful in conjunction with
+    :class:`torch.nn.parallel.DistributedDataParallel`. In such a case, each
+    process can pass a :class`~torch.utils.data.DistributedSampler` instance as a
+    :class:`~torch.utils.data.DataLoader` sampler, and load a subset of the
+    original dataset that is exclusive to it.
+
+    .. note::
+        Dataset is assumed to be of constant size.
+
+    Arguments:
+        dataset: Dataset used for sampling.
+        num_replicas (int, optional): Number of processes participating in
+            distributed training. By default, :attr:`rank` is retrieved from the
+            current distributed group.
+        rank (int, optional): Rank of the current process within :attr:`num_replicas`.
+            By default, :attr:`rank` is retrieved from the current distributed
+            group.
+        shuffle (bool, optional): If ``True`` (default), sampler will shuffle the
+            indices.
+        seed (int, optional): random seed used to shuffle the sampler if
+            :attr:`shuffle=True`. This number should be identical across all
+            processes in the distributed group. Default: ``0``.
+
+    .. warning::
+        In distributed mode, calling the :meth`set_epoch(epoch) <set_epoch>` method at
+        the beginning of each epoch **before** creating the :class:`DataLoader` iterator
+        is necessary to make shuffling work properly across multiple epochs. Otherwise,
+        the same ordering will be always used.
+
+    Example::
+
+        >>> sampler = DistributedSampler(dataset) if is_distributed else None
+        >>> loader = DataLoader(dataset, shuffle=(sampler is None),
+        ...                     sampler=sampler)
+        >>> for epoch in range(start_epoch, n_epochs):
+        ...     if is_distributed:
+        ...         sampler.set_epoch(epoch)
+        ...     train(loader)
+    """
+
+    def __init__(self, dataset, num_replicas=None, rank=None, shuffle=False, seed=0):
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = dist.get_rank()
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        # self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
+        # self.total_size = self.num_samples * self.num_replicas
+        self.total_size = len(self.dataset)         # true value without extra samples
+        indices = list(range(self.total_size))
+        indices = indices[self.rank:self.total_size:self.num_replicas]
+        self.num_samples = len(indices)             # true value without extra samples
+
+        self.shuffle = shuffle
+        self.seed = seed
+
+    def __iter__(self):
+        if self.shuffle:
+            # deterministically shuffle based on epoch and seed
+            g = torch.Generator()
+            g.manual_seed(self.seed + self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+        else:
+            indices = list(range(len(self.dataset)))
+
+
+        # # add extra samples to make it evenly divisible
+        # indices += indices[:(self.total_size - len(indices))]
+        # assert len(indices) == self.total_size
+
+        # subsample
+        indices = indices[self.rank:self.total_size:self.num_replicas]
+        assert len(indices) == self.num_samples
+
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+    def set_epoch(self, epoch):
+        r"""
+        Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas
+        use a different random ordering for each epoch. Otherwise, the next iteration of this
+        sampler will yield the same ordering.
+
+        Arguments:
+            epoch (int): _epoch number.
+        """
+        self.epoch = epoch

deblur/src/lambda_networks/__init__.py

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

deblur/src/lambda_networks/lambda_networks.py

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

deblur/src/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

deblur/src/launch.py

@@ -0,0 +1,55 @@
+""" distributed launcher adopted from torch.distributed.launch
+    usage example: https://github.com/facebookresearch/maskrcnn-benchmark
+    This enables using multiprocessing for each spawned process (as they are treated as main processes)
+"""
+import sys
+import subprocess
+from argparse import ArgumentParser, REMAINDER
+
+from utils import str2bool, int2str
+
+def parse_args():
+    parser = ArgumentParser(description="PyTorch distributed training launch "
+                                        "helper utilty that will spawn up "
+                                        "multiple distributed processes")
+
+
+    parser.add_argument('--n_GPUs', type=int, default=1, help='the number of GPUs for training')
+
+    # positional
+    parser.add_argument("training_script", type=str,
+                        help="The full path to the single GPU training "
+                             "program/script to be launched in parallel, "
+                             "followed by all the arguments for the "
+                             "training script")
+
+    # rest from the training program
+    parser.add_argument('training_script_args', nargs=REMAINDER)
+    return parser.parse_args()
+
+def main():
+    args = parse_args()
+
+    processes = []
+    for rank in range(0, args.n_GPUs):
+        cmd = [sys.executable]
+
+        cmd.append(args.training_script)
+        cmd.extend(args.training_script_args)
+
+        cmd += ['--distributed', 'True']
+        cmd += ['--launched', 'True']
+        cmd += ['--n_GPUs', str(args.n_GPUs)]
+        cmd += ['--rank', str(rank)]
+
+        process = subprocess.Popen(cmd)
+        processes.append(process)
+
+    for process in processes:
+        process.wait()
+        if process.returncode != 0:
+            raise subprocess.CalledProcessError(returncode=process.returncode,
+                                                cmd=cmd)
+
+if __name__ == "__main__":
+    main()

deblur/src/loss/__init__.py

@@ -0,0 +1,500 @@
+import os
+from importlib import import_module
+
+import torch
+from torch import nn
+import torch.distributed as dist
+
+import matplotlib.pyplot as plt
+plt.switch_backend('agg')	# https://github.com/matplotlib/matplotlib/issues/3466
+
+from .metric import PSNR, SSIM
+
+from utils import interact
+
+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(torch.nn.modules.loss._Loss):
+    def __init__(self, args, epoch=None, model=None, optimizer=None):
+        """
+            input:
+                args.loss       use '+' to sum over different loss functions
+                                use '*' to specify the loss weight
+
+                example:
+                    1*MSE+0.5*VGG54
+                                loss = sum of MSE and VGG54(weight=0.5)
+
+                args.measure    similar to args.loss, but without weight
+
+                example:
+                    MSE+PSNR
+                                measure MSE and PSNR, independently
+        """
+        super(Loss, self).__init__()
+
+        self.args = args
+
+        self.rgb_range = args.rgb_range
+        self.gamma=args.gamma
+        self.device_type = args.device_type
+        self.synchronized = False
+
+        self.epoch = args.start_epoch if epoch is None else epoch
+        self.save_dir = args.save_dir
+        self.save_name = os.path.join(self.save_dir, 'loss.pt')
+
+        # self.training = True
+        self.validating = False
+        self.testing = False
+        self.mode = 'train'
+        self.modes = ('train', 'val', 'test')
+
+        # Loss
+        self.loss = nn.ModuleDict()
+        self.loss_types = []
+        self.weight = {}
+        self.buffer=[0.0]*args.n_scales
+
+        self.loss_stat = {mode:{} for mode in self.modes}
+        # loss_stat[mode][loss_type][epoch] = loss_value
+        # loss_stat[mode]['Total'][epoch] = loss_total
+
+        for weighted_loss in args.loss.split('+'):
+            w, l = weighted_loss.split('*')
+            l = l.upper()
+            if l in ('ABS', 'L1'):
+                loss_type = 'L1'
+                func = nn.L1Loss()
+            elif l in ('MSE', 'L2'):
+                loss_type = 'L2'
+                func = nn.MSELoss()
+            elif l in ('ADV', 'GAN'):
+                loss_type = 'ADV'
+                m = import_module('loss.adversarial')
+                func = getattr(m, 'Adversarial')(args, model, optimizer)
+            else:
+                loss_type = l
+                m = import_module*'loss.{}'.format(l.lower())
+                func = getattr(m, l)(args)
+
+            self.loss_types += [loss_type]
+            self.loss[loss_type] = func
+            self.weight[loss_type] = float(w)
+
+        print('Loss function: {}'.format(args.loss))
+
+        # Metrics
+        self.do_measure = args.metric.lower() != 'none'
+
+        self.metric = nn.ModuleDict()
+        self.metric_types = []
+        self.metric_stat = {mode:{} for mode in self.modes}
+        # metric_stat[mode][metric_type][epoch] = metric_value
+
+        if self.do_measure:
+            for metric_type in args.metric.split(','):
+                metric_type = metric_type.upper()
+                if metric_type == 'PSNR':
+                    metric_func = PSNR()
+                elif metric_type == 'SSIM':
+                    metric_func = SSIM(args.device_type)    # single precision
+                else:
+                    raise NotImplementedError
+
+                self.metric_types += [metric_type]
+                self.metric[metric_type] = metric_func
+
+        print('Metrics: {}'.format(args.metric))
+
+        if args.start_epoch != 1:
+            self.load(args.start_epoch - 1)
+
+        for mode in self.modes:
+            for loss_type in self.loss:
+                if loss_type not in self.loss_stat[mode]:
+                    self.loss_stat[mode][loss_type] = {}   # initialize loss
+
+            if 'Total' not in self.loss_stat[mode]:
+                self.loss_stat[mode]['Total'] = {}
+
+            if self.do_measure:
+                for metric_type in self.metric:
+                    if metric_type not in self.metric_stat[mode]:
+                        self.metric_stat[mode][metric_type] = {}
+
+        self.count = 0
+        self.count_m = 0
+
+        self.to(args.device, dtype=args.dtype)
+
+    def train(self, mode=True):
+        super(Loss, self).train(mode)
+        if mode:
+            self.validating = False
+            self.testing = False
+            self.mode = 'train'
+        else:   # default test mode
+            self.validating = False
+            self.testing = True
+            self.mode = 'test'
+
+    def validate(self):
+        super(Loss, self).eval()
+        # self.training = False
+        self.validating = True
+        self.testing = False
+        self.mode = 'val'
+
+    def test(self):
+        super(Loss, self).eval()
+        # self.training = False
+        self.validating = False
+        self.testing = True
+        self.mode = 'test'
+
+    def forward(self, input, target):
+        self.synchronized = False
+
+        loss = 0
+        weights=[0.32,0.08,0.02,0.01,0.005]
+        if len(input)>len(weights):
+            for i in range(len(input)-len(weights)):
+                weights.append(weights[-1]*0.5)
+        if len(input)==1:
+             weights=[1.0]
+        weights=weights[::-1]
+        def _ms_forward(input, target, func):
+            if isinstance(input, (list, tuple)): # loss for list output
+                _loss,buffer_lst=sequence_loss(input,target[0],func,gamma=self.gamma)
+                # first=func(input[0],target[0])
+                # _loss = first*weights[0]
+                # self.buffer=[first.item()]
+                # for i in range(1,len(input)):
+                    # tmp=func(input[i],target[0])
+                    # self.buffer.append(tmp.item())
+                    # _loss+=tmp*weights[i]
+                self.buffer=buffer_lst
+                return _loss
+            elif isinstance(input, dict):   # loss for dict output
+                _loss = []
+                for key in input:
+                    _loss += [func(input[key], target[key])]
+                return sum(_loss)
+            else:   # loss for tensor output
+                return func(input, target)
+
+        # initialize
+        if self.count == 0:
+            for loss_type in self.loss_types:
+                self.loss_stat[self.mode][loss_type][self.epoch] = 0
+            self.loss_stat[self.mode]['Total'][self.epoch] = 0
+
+        if isinstance(input, list):
+            count = input[0].shape[0]
+        else:   # Tensor
+            count = input.shape[0]  # batch size
+
+        isnan = False
+        for loss_type in self.loss_types:
+
+            if loss_type == 'ADV':
+                _loss = self.loss[loss_type](input[0], target[0], self.training) * self.weight[loss_type]
+            else:
+                _loss = _ms_forward(input, target, self.loss[loss_type]) * self.weight[loss_type]
+
+            if torch.isnan(_loss):
+                isnan = True    # skip recording (will also be skipped at backprop)
+            else:
+                self.loss_stat[self.mode][loss_type][self.epoch] += _loss.item() * count
+                self.loss_stat[self.mode]['Total'][self.epoch] += _loss.item() * count
+
+            loss += _loss
+
+        if not isnan:
+            self.count += count
+
+        if not self.training and self.do_measure:
+            self.measure(input, target)
+
+        return loss
+
+    def measure(self, input, target):
+        if isinstance(input, (list, tuple)):
+            self.measure(input[0], target[0])
+            return
+        elif isinstance(input, dict):
+            first_key = list(input.keys())[0]
+            self.measure(input[first_key], target[first_key])
+            return
+        else:
+            pass
+
+        if self.count_m == 0:
+            for metric_type in self.metric_stat[self.mode]:
+                self.metric_stat[self.mode][metric_type][self.epoch] = 0
+
+        if isinstance(input, list):
+            count = input[0].shape[0]
+        else:   # Tensor
+            count = input.shape[0]  # batch size
+
+        for metric_type in self.metric_stat[self.mode]:
+
+            input = input.clamp(0, self.rgb_range)  # not in_place
+            if self.rgb_range==1:
+                input*=255
+                target*=255
+                input.round_()
+                target.round_()
+            if self.rgb_range == 255:
+                target*=255
+                input.round_()
+
+            _metric = self.metric[metric_type](input, target)
+            self.metric_stat[self.mode][metric_type][self.epoch] += _metric.item() * count
+
+        self.count_m += count
+
+        return
+
+    def normalize(self):
+        if self.args.distributed:
+            dist.barrier()
+            if not self.synchronized:
+                self.all_reduce()
+
+        if self.count > 0:
+            for loss_type in self.loss_stat[self.mode]: # including 'Total'
+                self.loss_stat[self.mode][loss_type][self.epoch] /= self.count
+            self.count = 0
+
+        if self.count_m > 0:
+            for metric_type in self.metric_stat[self.mode]:
+                self.metric_stat[self.mode][metric_type][self.epoch] /= self.count_m
+            self.count_m = 0
+
+        return
+
+    def all_reduce(self, epoch=None):
+        # synchronize loss for distributed GPU processes
+
+        if epoch is None:
+            epoch = self.epoch
+
+        def _reduce_value(value, ReduceOp=dist.ReduceOp.SUM):
+            value_tensor = torch.Tensor([value]).to(self.args.device, self.args.dtype, non_blocking=True)
+            dist.all_reduce(value_tensor, ReduceOp, async_op=False)
+            value = value_tensor.item()
+            del value_tensor
+
+            return value
+
+        dist.barrier()
+        if self.count > 0:  # I assume this should be true
+            self.count = _reduce_value(self.count, dist.ReduceOp.SUM)
+
+            for loss_type in self.loss_stat[self.mode]:
+                self.loss_stat[self.mode][loss_type][epoch] = _reduce_value(
+                    self.loss_stat[self.mode][loss_type][epoch],
+                    dist.ReduceOp.SUM
+                )
+
+        if self.count_m > 0:
+            self.count_m = _reduce_value(self.count_m, dist.ReduceOp.SUM)
+
+            for metric_type in self.metric_stat[self.mode]:
+                self.metric_stat[self.mode][metric_type][epoch] = _reduce_value(
+                    self.metric_stat[self.mode][metric_type][epoch],
+                    dist.ReduceOp.SUM
+                )
+
+        self.synchronized = True
+
+        return
+
+    def print_metrics(self):
+
+        print(self.get_metric_desc())
+        return
+
+    def get_last_loss(self):
+        return self.loss_stat[self.mode]['Total'][self.epoch]
+
+    def get_loss_desc(self):
+
+        if self.mode == 'train':
+            desc_prefix = 'Train'
+        elif self.mode == 'val':
+            desc_prefix = 'Validation'
+        else:
+            desc_prefix = 'Test'
+
+        loss = self.loss_stat[self.mode]['Total'][self.epoch]
+        if self.count > 0:
+            loss /= self.count
+        desc = '{} Loss: {:.1f}'.format(desc_prefix, loss)
+
+        if self.mode in ('val', 'test'):
+            metric_desc = self.get_metric_desc()
+            desc = '{}{}'.format(desc, metric_desc)
+
+        return desc
+
+    def get_metric_desc(self):
+        desc = ''
+        for metric_type in self.metric_stat[self.mode]:
+            measured = self.metric_stat[self.mode][metric_type][self.epoch]
+            if self.count_m > 0:
+                measured /= self.count_m
+
+            if metric_type == 'PSNR':
+                desc += ' {}: {:2.2f}'.format(metric_type, measured)
+            elif metric_type == 'SSIM':
+                desc += ' {}: {:1.4f}'.format(metric_type, measured)
+            else:
+                desc += ' {}: {:2.4f}'.format(metric_type, measured)
+
+        return desc
+
+    def step(self, plot_name=None):
+        self.normalize()
+        self.plot(plot_name)
+        if not self.training and self.do_measure:
+            # self.print_metrics()
+            self.plot_metric()
+        # self.epoch += 1
+
+        return
+
+    def save(self):
+
+        state = {
+            'loss_stat': self.loss_stat,
+            'metric_stat': self.metric_stat,
+        }
+        torch.save(state, self.save_name)
+
+        return
+
+    def load(self, epoch=None):
+
+        print('Loading loss record from {}'.format(self.save_name))
+        if os.path.exists(self.save_name):
+            state = torch.load(self.save_name, map_location=self.args.device)
+
+            self.loss_stat = state['loss_stat']
+            if 'metric_stat' in state:
+                self.metric_stat = state['metric_stat']
+            else:
+                pass
+        else:
+            print('no loss record found for {}!'.format(self.save_name))
+
+        if epoch is not None:
+            self.epoch = epoch
+
+        return
+
+    def plot(self, plot_name=None, metric=False):
+
+        self.plot_loss(plot_name)
+
+        if metric:
+            self.plot_metric(plot_name)
+        # else:
+        #     self.plot_loss(plot_name)
+
+        return
+
+
+    def plot_loss(self, plot_name=None):
+        if plot_name is None:
+            plot_name = os.path.join(self.save_dir, "{}_loss.pdf".format(self.mode))
+
+        title = "{} loss".format(self.mode)
+
+        fig = plt.figure()
+        plt.title(title)
+        plt.xlabel('epochs')
+        plt.ylabel('loss')
+        plt.grid(True, linestyle=':')
+
+        for loss_type, loss_record in self.loss_stat[self.mode].items(): # including Total
+            axis = sorted([epoch for epoch in loss_record.keys() if epoch <= self.epoch])
+            value = [self.loss_stat[self.mode][loss_type][epoch] for epoch in axis]
+            label = loss_type
+
+            plt.plot(axis, value, label=label)
+
+        plt.xlim(0, self.epoch)
+        plt.legend()
+        plt.savefig(plot_name)
+        plt.close(fig)
+
+        return
+
+    def plot_metric(self, plot_name=None):
+        # assume there are only max 2 metrics
+        if plot_name is None:
+            plot_name = os.path.join(self.save_dir, "{}_metric.pdf".format(self.mode))
+
+        title = "{} metrics".format(self.mode)
+
+        fig, ax1 = plt.subplots()
+        plt.title(title)
+        plt.grid(True, linestyle=':')
+        ax1.set_xlabel('epochs')
+
+        plots = None
+        for metric_type, metric_record in self.metric_stat[self.mode].items():
+            axis = sorted([epoch for epoch in metric_record.keys() if epoch <= self.epoch])
+            value = [metric_record[epoch] for epoch in axis]
+            label = metric_type
+
+            if metric_type == 'PSNR':
+                ax = ax1
+                color='C0'
+            elif metric_type == 'SSIM':
+                ax2 = ax1.twinx()
+                ax = ax2
+                color='C1'
+
+            ax.set_ylabel(metric_type)
+            if plots is None:
+                plots = ax.plot(axis, value, label=label, color=color)
+            else:
+                plots += ax.plot(axis, value, label=label, color=color)
+
+        labels = [plot.get_label() for plot in plots]
+        plt.legend(plots, labels)
+        plt.xlim(0, self.epoch)
+        plt.savefig(plot_name)
+        plt.close(fig)
+
+        return
+
+    def sort(self):
+        # sort the loss/metric record
+        for mode in self.modes:
+            for loss_type, loss_epochs in self.loss_stat[mode].items():
+                self.loss_stat[mode][loss_type] = {epoch: loss_epochs[epoch] for epoch in sorted(loss_epochs)}
+
+            for metric_type, metric_epochs in self.metric_stat[mode].items():
+                self.metric_stat[mode][metric_type] = {epoch: metric_epochs[epoch] for epoch in sorted(metric_epochs)}
+
+        return self

deblur/src/loss/adversarial.py

@@ -0,0 +1,52 @@
+import torch
+import torch.nn as nn
+
+from utils import interact
+
+import torch.cuda.amp as amp
+
+class Adversarial(nn.modules.loss._Loss):
+    # pure loss function without saving & loading option
+    # but trains deiscriminator
+    def __init__(self, args, model, optimizer):
+        super(Adversarial, self).__init__()
+        self.args = args
+        self.model = model.model
+        self.optimizer = optimizer
+        self.scaler = amp.GradScaler(
+            init_scale=self.args.init_scale,
+            enabled=self.args.amp
+        )
+
+        self.gan_k = 1
+
+        self.BCELoss = nn.BCEWithLogitsLoss()
+
+    def forward(self, fake, real, training=False):
+        if training:
+            # update discriminator
+            fake_detach = fake.detach()
+            for _ in range(self.gan_k):
+                self.optimizer.D.zero_grad()
+                # d: B x 1 tensor
+                with amp.autocast(self.args.amp):
+                    d_fake = self.model.D(fake_detach)
+                    d_real = self.model.D(real)
+
+                    label_fake = torch.zeros_like(d_fake)
+                    label_real = torch.ones_like(d_real)
+
+                    loss_d = self.BCELoss(d_fake, label_fake) + self.BCELoss(d_real, label_real)
+
+                self.scaler.scale(loss_d).backward(retain_graph=False)
+                self.scaler.step(self.optimizer.D)
+                self.scaler.update()
+        else:
+            d_real = self.model.D(real)
+            label_real = torch.ones_like(d_real)
+
+        # update generator (outside here)
+        d_fake_bp = self.model.D(fake)
+        loss_g = self.BCELoss(d_fake_bp, label_real)
+
+        return loss_g

deblur/src/loss/metric.py

@@ -0,0 +1,112 @@
+# from skimage.metrics import peak_signal_noise_ratio, structural_similarity
+
+import torch
+from torch import nn
+
+def _expand(img):
+    if img.ndim < 4:
+        img = img.expand([1] * (4-img.ndim) + list(img.shape))
+
+    return img
+
+class PSNR(nn.Module):
+    def __init__(self):
+        super(PSNR, self).__init__()
+
+    def forward(self, im1, im2, data_range=None):
+        # tensor input, constant output
+
+        if data_range is None:
+            data_range = 255 if im1.max() > 1 else 1
+
+        se = (im1-im2)**2
+        se = _expand(se)
+
+        mse = se.mean(dim=list(range(1, se.ndim)))
+        psnr = 10 * (data_range**2/mse).log10().mean()
+
+        return psnr
+
+class SSIM(nn.Module):
+    def __init__(self, device_type='cpu', dtype=torch.float32):
+        super(SSIM, self).__init__()
+
+        self.device_type = device_type
+        self.dtype = dtype      # SSIM in half precision could be inaccurate
+
+        def _get_ssim_weight():
+            truncate = 3.5
+            sigma = 1.5
+            r = int(truncate * sigma + 0.5)  # radius as in ndimage
+            win_size = 2 * r + 1
+            nch = 3
+
+            weight = torch.Tensor([-(x - win_size//2)**2/float(2*sigma**2) for x in range(win_size)]).exp().unsqueeze(1)
+            weight = weight.mm(weight.t())
+            weight /= weight.sum()
+            weight = weight.repeat(nch, 1, 1, 1)
+
+            return weight
+
+        self.weight = _get_ssim_weight().to(self.device_type, dtype=self.dtype, non_blocking=True)
+
+    def forward(self, im1, im2, data_range=None):
+        """Implementation adopted from skimage.metrics.structural_similarity
+        Default arguments set to multichannel=True, gaussian_weight=True, use_sample_covariance=False
+        """
+
+        im1 = im1.to(self.device_type, dtype=self.dtype, non_blocking=True)
+        im2 = im2.to(self.device_type, dtype=self.dtype, non_blocking=True)
+
+        K1 = 0.01
+        K2 = 0.03
+        sigma = 1.5
+
+        truncate = 3.5
+        r = int(truncate * sigma + 0.5)  # radius as in ndimage
+        win_size = 2 * r + 1
+
+        im1 = _expand(im1)
+        im2 = _expand(im2)
+
+        nch = im1.shape[1]
+
+        if im1.shape[2] < win_size or im1.shape[3] < win_size:
+            raise ValueError(
+                "win_size exceeds image extent.  If the input is a multichannel "
+                "(color) image, set multichannel=True.")
+
+        if data_range is None:
+            data_range = 255 if im1.max() > 1 else 1
+
+        def filter_func(img):   # no padding
+            return nn.functional.conv2d(img, self.weight, groups=nch).to(self.dtype)
+            # return torch.conv2d(img, self.weight, groups=nch).to(self.dtype)
+
+        # compute (weighted) means
+        ux = filter_func(im1)
+        uy = filter_func(im2)
+
+        # compute (weighted) variances and covariances
+        uxx = filter_func(im1 * im1)
+        uyy = filter_func(im2 * im2)
+        uxy = filter_func(im1 * im2)
+        vx = (uxx - ux * ux)
+        vy = (uyy - uy * uy)
+        vxy = (uxy - ux * uy)
+
+        R = data_range
+        C1 = (K1 * R) ** 2
+        C2 = (K2 * R) ** 2
+
+        A1, A2, B1, B2 = ((2 * ux * uy + C1,
+                        2 * vxy + C2,
+                        ux ** 2 + uy ** 2 + C1,
+                        vx + vy + C2))
+        D = B1 * B2
+        S = (A1 * A2) / D
+
+        # compute (weighted) mean of ssim
+        mssim = S.mean()
+
+        return mssim

deblur/src/main.py

@@ -0,0 +1,67 @@
+"""main file that does everything"""
+from utils import interact
+
+from option import args, setup, cleanup
+from data import Data
+from model import Model
+from loss import Loss
+from optim import Optimizer
+from train import Trainer
+
+def main_worker(rank, args):
+    args.rank = rank
+    args = setup(args)
+
+    loaders = Data(args).get_loader()
+    model = Model(args)
+    model.parallelize()
+    optimizer = Optimizer(args, model)
+
+    criterion = Loss(args, model=model, optimizer=optimizer)
+
+    trainer = Trainer(args, model, criterion, optimizer, loaders)
+
+    if args.stay:
+        interact(local=locals())
+        exit()
+
+    if args.demo:
+        trainer.evaluate(epoch=args.start_epoch, mode='demo')
+        exit()
+
+    for epoch in range(1, args.start_epoch):
+        if args.do_validate:
+            if epoch % args.validate_every == 0:
+                trainer.fill_evaluation(epoch, 'val')
+        if args.do_test:
+            if epoch % args.test_every == 0:
+                trainer.fill_evaluation(epoch, 'test')
+
+    for epoch in range(args.start_epoch, args.end_epoch+1):
+        if args.do_train:
+            trainer.train(epoch)
+
+        if args.do_validate:
+            if epoch % args.validate_every == 0:
+                if trainer.epoch != epoch:
+                    trainer.load(epoch)
+                trainer.validate(epoch)
+
+        if args.do_test:
+            if epoch % args.test_every == 0:
+                if trainer.epoch != epoch:
+                    trainer.load(epoch)
+                trainer.test(epoch)
+
+        if args.rank == 0 or not args.launched:
+            print('')
+
+    trainer.imsaver.join_background()
+
+    cleanup(args)
+
+def main():
+    main_worker(args.rank, args)
+
+if __name__ == "__main__":
+    main()

deblur/src/model/LamResNet.py

@@ -0,0 +1,72 @@
+import torch.nn as nn
+
+from . import common
+
+from lambda_networks import LambdaLayer
+
+
+def build_model(args):
+    return ResNet(args)
+
+
+class ResNet(nn.Module):
+    def __init__(
+        self,
+        args,
+        in_channels=3,
+        out_channels=3,
+        n_feats=None,
+        kernel_size=None,
+        n_resblocks=None,
+        mean_shift=True,
+    ):
+        super(ResNet, self).__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.n_feats = args.n_feats if n_feats is None else n_feats
+        self.kernel_size = args.kernel_size if kernel_size is None else kernel_size
+        self.n_resblocks = args.n_resblocks if n_resblocks is None else n_resblocks
+
+        self.mean_shift = mean_shift
+        self.rgb_range = args.rgb_range
+        self.mean = self.rgb_range / 2
+
+        modules = []
+        modules.append(
+            common.default_conv(self.in_channels, self.n_feats, self.kernel_size)
+        )
+        for _ in range(self.n_resblocks // 3):
+            modules.append(common.ResBlock(self.n_feats, self.kernel_size))
+        modules.append(
+            LambdaLayer(
+                dim=self.n_feats, dim_out=self.n_feats, r=23, dim_k=16, heads=4, dim_u=1
+            )
+        )
+        for _ in range(self.n_resblocks // 3):
+            modules.append(common.ResBlock(self.n_feats, self.kernel_size))
+        modules.append(
+            LambdaLayer(
+                dim=self.n_feats, dim_out=self.n_feats, r=7, dim_k=16, heads=4, dim_u=4
+            )
+        )
+        for _ in range(self.n_resblocks // 3):
+            modules.append(common.ResBlock(self.n_feats, self.kernel_size))
+        modules.append(
+            common.default_conv(self.n_feats, self.n_feats, self.kernel_size)
+        )
+        modules.append(common.default_conv(self.n_feats, self.out_channels, 1))
+
+        self.body = nn.Sequential(*modules)
+
+    def forward(self, input):
+        if self.mean_shift:
+            input = input - self.mean
+
+        output = self.body(input)
+
+        if self.mean_shift:
+            output = output + self.mean
+
+        return output

deblur/src/model/MSResNet.py

@@ -0,0 +1,67 @@
+import torch
+import torch.nn as nn
+
+from . import common
+from .ResNet import ResNet
+
+
+def build_model(args):
+    return MSResNet(args)
+
+class conv_end(nn.Module):
+    def __init__(self, in_channels=3, out_channels=3, kernel_size=5, ratio=2):
+        super(conv_end, self).__init__()
+
+        modules = [
+            common.default_conv(in_channels, out_channels, kernel_size),
+            nn.PixelShuffle(ratio)
+        ]
+
+        self.uppath = nn.Sequential(*modules)
+
+    def forward(self, x):
+        return self.uppath(x)
+
+class MSResNet(nn.Module):
+    def __init__(self, args):
+        super(MSResNet, self).__init__()
+
+        self.rgb_range = args.rgb_range
+        self.mean = self.rgb_range / 2
+
+        self.n_resblocks = args.n_resblocks
+        self.n_feats = args.n_feats
+        self.kernel_size = args.kernel_size
+
+        self.n_scales = args.n_scales
+
+        self.body_models = nn.ModuleList([
+            ResNet(args, 3, 3, mean_shift=False),
+        ])
+        for _ in range(1, self.n_scales):
+            self.body_models.insert(0, ResNet(args, 6, 3, mean_shift=False))
+
+        self.conv_end_models = nn.ModuleList([None])
+        for _ in range(1, self.n_scales):
+            self.conv_end_models += [conv_end(3, 12)]
+
+    def forward(self, input_pyramid):
+
+        scales = range(self.n_scales-1, -1, -1)    # 0: fine, 2: coarse
+
+        for s in scales:
+            input_pyramid[s] = input_pyramid[s] - self.mean
+
+        output_pyramid = [None] * self.n_scales
+
+        input_s = input_pyramid[-1]
+        for s in scales:    # [2, 1, 0]
+            output_pyramid[s] = self.body_models[s](input_s)
+            if s > 0:
+                up_feat = self.conv_end_models[s](output_pyramid[s])
+                input_s = torch.cat((input_pyramid[s-1], up_feat), 1)
+
+        for s in scales:
+            output_pyramid[s] = output_pyramid[s] + self.mean
+
+        return output_pyramid

deblur/src/model/MSResNetLambda.py

@@ -0,0 +1,77 @@
+import torch
+import torch.nn as nn
+
+from . import common
+from .ResNet import ResNet
+from lambda_network import LambdaLayer
+
+def build_model(args):
+    return MSResNet(args)
+
+class conv_end(nn.Module):
+    def __init__(self, in_channels=3, out_channels=3, kernel_size=5, ratio=2):
+        super(conv_end, self).__init__()
+
+        modules = [
+            common.default_conv(in_channels, out_channels, kernel_size),
+            nn.PixelShuffle(ratio)
+        ]
+
+        self.uppath = nn.Sequential(*modules)
+
+    def forward(self, x):
+        return self.uppath(x)
+
+class MSResNet(nn.Module):
+    def __init__(self, args):
+        super(MSResNet, self).__init__()
+
+        self.rgb_range = args.rgb_range
+        self.mean = self.rgb_range / 2
+
+        self.n_resblocks = args.n_resblocks
+        self.n_feats = args.n_feats
+        self.kernel_size = args.kernel_size
+
+        self.n_scales = args.n_scales
+
+        self.body_models = nn.ModuleList([
+            ResNet(args, 3, 3, mean_shift=False),
+        ])
+        self.lambda_models = nn.ModuleList([
+            LambdaLayer(
+    dim = 32,       # channels going in
+    dim_out = 32,   # channels out
+    n = 64 * 64,    # number of input pixels (64 x 64 image)
+    dim_k = 16,     # key dimension
+    heads = 4,      # number of heads, for multi-query
+    dim_u = 1       # 'intra-depth' dimension
+)
+        ])
+        for _ in range(1, self.n_scales):
+            self.body_models.insert(0, ResNet(args, 6, 3, mean_shift=False))
+
+        self.conv_end_models = nn.ModuleList([None])
+        for _ in range(1, self.n_scales):
+            self.conv_end_models += [conv_end(3, 12)]
+
+    def forward(self, input_pyramid):
+
+        scales = range(self.n_scales-1, -1, -1)    # 0: fine, 2: coarse
+
+        for s in scales:
+            input_pyramid[s] = input_pyramid[s] - self.mean
+
+        output_pyramid = [None] * self.n_scales
+
+        input_s = input_pyramid[-1]
+        for s in scales:    # [2, 1, 0]
+            output_pyramid[s] = self.body_models[s](input_s)
+            if s > 0:
+                up_feat = self.conv_end_models[s](output_pyramid[s])
+                input_s = torch.cat((input_pyramid[s-1], up_feat), 1)
+
+        for s in scales:
+            output_pyramid[s] = output_pyramid[s] + self.mean
+
+        return output_pyramid

deblur/src/model/RaftNet.py

@@ -0,0 +1,120 @@
+import torch.nn as nn
+
+import torch
+from . import common
+
+from lambda_networks import LambdaLayer
+
+
+def build_model(args):
+    return ResNet(args)
+
+
+class ConvGRU(nn.Module):
+    def __init__(self, hidden_dim=128, input_dim=192+128):
+        super(ConvGRU, self).__init__()
+        self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1)
+        self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1)
+        self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1)
+
+    def forward(self, h, x):
+        hx = torch.cat([h, x], dim=1)
+
+        z = torch.sigmoid(self.convz(hx))
+        r = torch.sigmoid(self.convr(hx))
+        q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)))
+
+        # h = (1-z) * h + z * q
+        # return h
+        return (1-z) * h + z * q
+
+class SepConvGRU(nn.Module):
+    def __init__(self, hidden_dim=128, input_dim=192+128):
+        super(SepConvGRU, self).__init__()
+        self.convz1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
+        self.convr1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
+        self.convq1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
+
+        self.convz2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
+        self.convr2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
+        self.convq2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
+
+
+    def forward(self, h, x):
+        # horizontal
+        hx = torch.cat([h, x], dim=1)
+        z = torch.sigmoid(self.convz1(hx))
+        r = torch.sigmoid(self.convr1(hx))
+        q = torch.tanh(self.convq1(torch.cat([r*h, x], dim=1)))        
+        h = (1-z) * h + z * q
+
+        # vertical
+        hx = torch.cat([h, x], dim=1)
+        z = torch.sigmoid(self.convz2(hx))
+        r = torch.sigmoid(self.convr2(hx))
+        q = torch.tanh(self.convq2(torch.cat([r*h, x], dim=1)))       
+        h = (1-z) * h + z * q
+
+        return h
+
+
+class ResNet(nn.Module):
+    def __init__(
+        self,
+        args,
+        in_channels=3,
+        out_channels=3,
+        n_feats=None,
+        kernel_size=None,
+        n_resblocks=None,
+        mean_shift=True,
+    ):
+        super(ResNet, self).__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.n_feats = args.n_feats if n_feats is None else n_feats
+        self.kernel_size = args.kernel_size if kernel_size is None else kernel_size
+        self.n_resblocks = args.n_resblocks if n_resblocks is None else n_resblocks
+
+        self.mean_shift = mean_shift
+        self.rgb_range = args.rgb_range
+        self.mean = self.rgb_range / 2
+
+        modules = []
+        m_head=[common.default_conv(self.in_channels, self.n_feats, self.kernel_size)]
+        for i in range(3):
+            m_head.append(common.ResBlock(self.n_feats, self.kernel_size))
+        for _ in range(self.n_resblocks // 2):
+            modules.append(common.ResBlock(self.n_feats, self.kernel_size))
+        modules.append(
+            LambdaLayer(
+                dim=self.n_feats, dim_out=self.n_feats, r=23, dim_k=16, heads=4, dim_u=4
+            )
+        )
+        for _ in range(self.n_resblocks // 2):
+            modules.append(common.ResBlock(self.n_feats, self.kernel_size))
+        m_tail=[]
+        
+        for i in range(3):
+            m_tail.append(common.ResBlock(self.n_feats, self.kernel_size))
+        
+        m_tail.append(
+            common.default_conv(self.n_feats, self.out_channels, self.kernel_size)
+        )
+        self.head=nn.Sequential(*m_head)
+        self.body = nn.Sequential(*modules)
+        self.tail=nn.Sequential(*m_tail)
+        self.gru=SepConvGRU(hidden_dim=self.n_feats,input_dim=self.n_feats)
+
+    def forward(self, input):
+        if self.mean_shift:
+            input = input - self.mean
+
+        output = self.body(input)
+
+        if self.mean_shift:
+            output = output + self.mean
+
+        return output

deblur/src/model/RecLamResNet.py

@@ -0,0 +1,54 @@
+import torch
+import torch.nn as nn
+
+from . import common
+from .LamResNet import ResNet
+
+
+def build_model(args):
+    return RecLamResNet(args)
+
+
+class conv_end(nn.Module):
+    def __init__(self, in_channels=3, out_channels=3, kernel_size=5, ratio=2):
+        super(conv_end, self).__init__()
+
+        modules = [
+            common.default_conv(in_channels, out_channels, kernel_size),
+            nn.PixelShuffle(ratio),
+        ]
+
+        self.uppath = nn.Sequential(*modules)
+
+    def forward(self, x):
+        return self.uppath(x)
+
+
+class RecLamResNet(nn.Module):
+    def __init__(self, args):
+        super(RecLamResNet, self).__init__()
+
+        self.rgb_range = args.rgb_range
+        self.mean = self.rgb_range / 2
+        self.is_detach=args.detach
+
+        self.n_resblocks = args.n_resblocks
+        self.n_feats = args.n_feats
+        self.kernel_size = args.kernel_size
+
+        self.n_scales = args.n_scales
+
+        self.body_model = ResNet(args, 3, 3, mean_shift=False)
+
+    def forward(self, input_lst):
+        # we use a reversed list for better compact
+        input_lst[0] = input_lst[0] - self.mean
+        output_lst = [None] * self.n_scales
+        last_output = input_lst[0]
+        for i in range(self.n_scales):
+            if self.is_detach:
+                last_output=last_output.detach()
+            output = self.body_model(last_output) + last_output
+            output_lst[self.n_scales-i-1] = output + self.mean
+            last_output = output
+        return output_lst

deblur/src/model/ResNet.py

@@ -0,0 +1,41 @@
+import torch.nn as nn
+
+from . import common
+
+def build_model(args):
+    return ResNet(args)
+
+class ResNet(nn.Module):
+    def __init__(self, args, in_channels=3, out_channels=3, n_feats=None, kernel_size=None, n_resblocks=None, mean_shift=True):
+        super(ResNet, self).__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.n_feats = args.n_feats if n_feats is None else n_feats
+        self.kernel_size = args.kernel_size if kernel_size is None else kernel_size
+        self.n_resblocks = args.n_resblocks if n_resblocks is None else n_resblocks
+
+        self.mean_shift = mean_shift
+        self.rgb_range = args.rgb_range
+        self.mean = self.rgb_range / 2
+
+        modules = []
+        modules.append(common.default_conv(self.in_channels, self.n_feats, self.kernel_size))
+        for _ in range(self.n_resblocks):
+            modules.append(common.ResBlock(self.n_feats, self.kernel_size))
+        modules.append(common.default_conv(self.n_feats, self.out_channels, self.kernel_size))
+
+        self.body = nn.Sequential(*modules)
+
+    def forward(self, input):
+        if self.mean_shift:
+            input = input - self.mean
+
+        output = self.body(input)
+
+        if self.mean_shift:
+            output = output + self.mean
+
+        return output
+

deblur/src/model/__init__.py

@@ -0,0 +1,136 @@
+import os
+import re
+from importlib import import_module
+
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DataParallel, DistributedDataParallel
+
+import torch.distributed as dist
+from torch.nn.utils import parameters_to_vector, vector_to_parameters
+
+from .discriminator import Discriminator
+
+from utils import interact
+
+class Model(nn.Module):
+    def __init__(self, args):
+        super(Model, self).__init__()
+
+        self.args = args
+        self.device = args.device
+        self.n_GPUs = args.n_GPUs
+        self.save_dir = os.path.join(args.save_dir, 'models')
+        os.makedirs(self.save_dir, exist_ok=True)
+
+        module = import_module('model.' + args.model)
+
+        self.model = nn.ModuleDict()
+        self.model.G = module.build_model(args)
+        if self.args.loss.lower().find('adv') >= 0:
+            self.model.D = Discriminator(self.args)
+        else:
+            self.model.D = None
+
+        self.to(args.device, dtype=args.dtype, non_blocking=True)
+        self.load(args.load_epoch, path=args.pretrained)
+
+    def parallelize(self):
+        if self.args.device_type == 'cuda':
+            if self.args.distributed:
+                Parallel = DistributedDataParallel
+                parallel_args = {
+                    "device_ids": [self.args.rank],
+                    "output_device": self.args.rank,
+                }
+            else:
+                Parallel = DataParallel
+                parallel_args = {
+                    'device_ids': list(range(self.n_GPUs)),
+                    'output_device': self.args.rank # always 0
+                }
+
+            for model_key in self.model:
+                if self.model[model_key] is not None:
+                    self.model[model_key] = Parallel(self.model[model_key], **parallel_args)
+
+    def forward(self, input):
+        return self.model.G(input)
+
+    def _save_path(self, epoch):
+        model_path = os.path.join(self.save_dir, 'model-{:d}.pt'.format(epoch))
+        return model_path
+
+    def state_dict(self):
+        state_dict = {}
+        for model_key in self.model:
+            if self.model[model_key] is not None:
+                parallelized = isinstance(self.model[model_key], (DataParallel, DistributedDataParallel))
+                if parallelized:
+                    state_dict[model_key] = self.model[model_key].module.state_dict()
+                else:
+                    state_dict[model_key] = self.model[model_key].state_dict()
+
+        return state_dict
+
+    def load_state_dict(self, state_dict, strict=True):
+        for model_key in self.model:
+            parallelized = isinstance(self.model[model_key], (DataParallel, DistributedDataParallel))
+            if model_key in state_dict:
+                if parallelized:
+                    self.model[model_key].module.load_state_dict(state_dict[model_key], strict)
+                else:
+                    self.model[model_key].load_state_dict(state_dict[model_key], strict)
+
+    def save(self, epoch):
+        torch.save(self.state_dict(), self._save_path(epoch))
+
+    def load(self, epoch=None, path=None):
+        if path:
+            model_name = path
+        elif isinstance(epoch, int):
+            if epoch < 0:
+                epoch = self.get_last_epoch()
+            if epoch == 0:   # epoch 0
+                # make sure model parameters are synchronized at initial
+                # for multi-node training (not in current implementation)
+                # self.synchronize()
+
+                return  # leave model as initialized
+
+            model_name = self._save_path(epoch)
+        else:
+            raise Exception('no epoch number or model path specified!')
+
+        print('Loading model from {}'.format(model_name))
+        state_dict = torch.load(model_name, map_location=self.args.device)
+        self.load_state_dict(state_dict)
+
+        return
+
+    def synchronize(self):
+        if self.args.distributed:
+            # synchronize model parameters across nodes
+            vector = parameters_to_vector(self.parameters())
+
+            dist.broadcast(vector, 0)   # broadcast parameters to other processes
+            if self.args.rank != 0:
+                vector_to_parameters(vector, self.parameters())
+
+            del vector
+
+        return
+
+    def get_last_epoch(self):
+        model_list = sorted(os.listdir(self.save_dir))
+        if len(model_list) == 0:
+            epoch = 0
+        else:
+            epoch = int(re.findall('\\d+', model_list[-1])[0]) # model example name model-100.pt
+
+        return epoch
+
+    def print(self):
+        print(self.model)
+
+        return

deblur/src/model/common.py

@@ -0,0 +1,161 @@
+import math
+
+import torch
+import torch.nn as nn
+
+def default_conv(in_channels, out_channels, kernel_size, bias=True, groups=1):
+    return nn.Conv2d(
+        in_channels, out_channels, kernel_size,
+        padding=(kernel_size // 2), bias=bias, groups=groups)
+
+def default_norm(n_feats):
+    return nn.BatchNorm2d(n_feats)
+
+def default_act():
+    return nn.ReLU(True)
+
+def empty_h(x, n_feats):
+    '''
+        create an empty hidden state
+
+        input
+            x:      B x T x 3 x H x W
+
+        output
+            h:      B x C x H/4 x W/4
+    '''
+    b = x.size(0)
+    h, w = x.size()[-2:]
+    return x.new_zeros((b, n_feats, h//4, w//4))
+
+class Normalization(nn.Conv2d):
+    """Normalize input tensor value with convolutional layer"""
+    def __init__(self, mean=(0, 0, 0), std=(1, 1, 1)):
+        super(Normalization, self).__init__(3, 3, kernel_size=1)
+        tensor_mean = torch.Tensor(mean)
+        tensor_inv_std = torch.Tensor(std).reciprocal()
+
+        self.weight.data = torch.eye(3).mul(tensor_inv_std).view(3, 3, 1, 1)
+        self.bias.data = torch.Tensor(-tensor_mean.mul(tensor_inv_std))
+
+        for params in self.parameters():
+            params.requires_grad = False
+
+class BasicBlock(nn.Sequential):
+    """Convolution layer + Activation layer"""
+    def __init__(
+        self, in_channels, out_channels, kernel_size, bias=True,
+        conv=default_conv, norm=False, act=default_act):
+
+        modules = []
+        modules.append(
+            conv(in_channels, out_channels, kernel_size, bias=bias))
+        if norm: modules.append(norm(out_channels))
+        if act: modules.append(act())
+
+        super(BasicBlock, self).__init__(*modules)
+
+class ResBlock(nn.Module):
+    def __init__(
+        self, n_feats, kernel_size, bias=True,
+        conv=default_conv, norm=False, act=default_act):
+
+        super(ResBlock, self).__init__()
+
+        modules = []
+        for i in range(2):
+            modules.append(conv(n_feats, n_feats, kernel_size, bias=bias))
+            if norm: modules.append(norm(n_feats))
+            if act and i == 0: modules.append(act())
+
+        self.body = nn.Sequential(*modules)
+
+    def forward(self, x):
+        res = self.body(x)
+        res += x
+
+        return res
+
+class ResBlock_mobile(nn.Module):
+    def __init__(
+        self, n_feats, kernel_size, bias=True,
+        conv=default_conv, norm=False, act=default_act, dropout=False):
+
+        super(ResBlock_mobile, self).__init__()
+
+        modules = []
+        for i in range(2):
+            modules.append(conv(n_feats, n_feats, kernel_size, bias=False, groups=n_feats))
+            modules.append(conv(n_feats, n_feats, 1, bias=False))
+            if dropout and i == 0: modules.append(nn.Dropout2d(dropout))
+            if norm: modules.append(norm(n_feats))
+            if act and i == 0: modules.append(act())
+
+        self.body = nn.Sequential(*modules)
+
+    def forward(self, x):
+        res = self.body(x)
+        res += x
+
+        return res
+
+class Upsampler(nn.Sequential):
+    def __init__(
+        self, scale, n_feats, bias=True,
+        conv=default_conv, norm=False, act=False):
+
+        modules = []
+        if (scale & (scale - 1)) == 0:    # Is scale = 2^n?
+            for _ in range(int(math.log(scale, 2))):
+                modules.append(conv(n_feats, 4 * n_feats, 3, bias))
+                modules.append(nn.PixelShuffle(2))
+                if norm: modules.append(norm(n_feats))
+                if act: modules.append(act())
+        elif scale == 3:
+            modules.append(conv(n_feats, 9 * n_feats, 3, bias))
+            modules.append(nn.PixelShuffle(3))
+            if norm: modules.append(norm(n_feats))
+            if act: modules.append(act())
+        else:
+            raise NotImplementedError
+
+        super(Upsampler, self).__init__(*modules)
+
+# Only support 1 / 2
+class PixelSort(nn.Module):
+    """The inverse operation of PixelShuffle
+    Reduces the spatial resolution, increasing the number of channels.
+    Currently, scale 0.5 is supported only.
+    Later, torch.nn.functional.pixel_sort may be implemented.
+    Reference:
+        http://pytorch.org/docs/0.3.0/_modules/torch/nn/modules/pixelshuffle.html#PixelShuffle
+        http://pytorch.org/docs/0.3.0/_modules/torch/nn/functional.html#pixel_shuffle
+    """
+    def __init__(self, upscale_factor=0.5):
+        super(PixelSort, self).__init__()
+        self.upscale_factor = upscale_factor
+
+    def forward(self, x):
+        b, c, h, w = x.size()
+        x = x.view(b, c, 2, 2, h // 2, w // 2)
+        x = x.permute(0, 1, 5, 3, 2, 4).contiguous()
+        x = x.view(b, 4 * c, h // 2, w // 2)
+
+        return x
+
+class Downsampler(nn.Sequential):
+    def __init__(
+        self, scale, n_feats, bias=True,
+        conv=default_conv, norm=False, act=False):
+
+        modules = []
+        if scale == 0.5:
+            modules.append(PixelSort())
+            modules.append(conv(4 * n_feats, n_feats, 3, bias))
+            if norm: modules.append(norm(n_feats))
+            if act: modules.append(act())
+        else:
+            raise NotImplementedError
+
+        super(Downsampler, self).__init__(*modules)
+

deblur/src/model/discriminator.py

@@ -0,0 +1,41 @@
+import torch.nn as nn
+
+class Discriminator(nn.Module):
+    def __init__(self, args):
+        super(Discriminator, self).__init__()
+
+        # self.args = args
+        n_feats = args.n_feats
+        kernel_size = args.kernel_size
+
+        def conv(kernel_size, in_channel, n_feats, stride, pad=None):
+            if pad is None:
+                pad = (kernel_size-1)//2
+
+            return nn.Conv2d(in_channel, n_feats, kernel_size, stride=stride, padding=pad, bias=False)
+
+        self.conv_layers = nn.ModuleList([
+            conv(kernel_size, 3,         n_feats//2, 1),    # 256
+            conv(kernel_size, n_feats//2, n_feats//2, 2),   # 128
+            conv(kernel_size, n_feats//2, n_feats,   1),
+            conv(kernel_size, n_feats,   n_feats,   2),     # 64
+            conv(kernel_size, n_feats,   n_feats*2, 1),
+            conv(kernel_size, n_feats*2, n_feats*2, 4),     # 16
+            conv(kernel_size, n_feats*2, n_feats*4, 1),
+            conv(kernel_size, n_feats*4, n_feats*4, 4),     # 4
+            conv(kernel_size, n_feats*4, n_feats*8, 1),
+            conv(4,           n_feats*8, n_feats*8, 4, 0),  # 1
+        ])
+
+        self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+        self.dense = nn.Conv2d(n_feats*8, 1, 1, bias=False)
+
+    def forward(self, x):
+
+        for layer in self.conv_layers:
+            x = self.act(layer(x))
+
+        x = self.dense(x)
+
+        return x
+

deblur/src/model/structure.py

@@ -0,0 +1,56 @@
+import torch.nn as nn
+
+from .common import ResBlock, default_conv
+
+def encoder(in_channels, n_feats):
+    """RGB / IR feature encoder
+    """
+
+    # in_channels == 1 or 3 or 4 or ....
+    # After 1st conv, B x n_feats x H x W
+    # After 2nd conv, B x 2n_feats x H/2 x W/2
+    # After 3rd conv, B x 3n_feats x H/4 x W/4
+    return nn.Sequential(
+        nn.Conv2d(in_channels, 1 * n_feats, 5, stride=1, padding=2),
+        nn.Conv2d(1 * n_feats, 2 * n_feats, 5, stride=2, padding=2),
+        nn.Conv2d(2 * n_feats, 3 * n_feats, 5, stride=2, padding=2),
+    )
+
+def decoder(out_channels, n_feats):
+    """RGB / IR / Depth decoder
+    """
+    # After 1st deconv, B x 2n_feats x H/2 x W/2
+    # After 2nd deconv, B x n_feats x H x W
+    # After 3rd conv, B x out_channels x H x W
+    deconv_kargs = {'stride': 2, 'padding': 1, 'output_padding': 1}
+
+    return nn.Sequential(
+        nn.ConvTranspose2d(3 * n_feats, 2 * n_feats, 3, **deconv_kargs),
+        nn.ConvTranspose2d(2 * n_feats, 1 * n_feats, 3, **deconv_kargs),
+        nn.Conv2d(n_feats, out_channels, 5, stride=1, padding=2),
+    )
+
+# def ResNet(n_feats, in_channels=None, out_channels=None):
+def ResNet(n_feats, kernel_size, n_blocks, in_channels=None, out_channels=None):
+    """sequential ResNet
+    """
+
+    # if in_channels is None:
+    #     in_channels = n_feats
+    # if out_channels is None:
+    #     out_channels = n_feats
+    # # currently not implemented
+
+    m = []
+
+    if in_channels is not None:
+        m += [default_conv(in_channels, n_feats, kernel_size)]
+
+    m += [ResBlock(n_feats, 3)] * n_blocks
+
+    if out_channels is not None:
+        m += [default_conv(n_feats, out_channels, kernel_size)]
+
+
+    return nn.Sequential(*m)
+

deblur/src/optim/__init__.py

@@ -0,0 +1,206 @@
+import torch
+import torch.optim as optim
+import torch.optim.lr_scheduler as lrs
+
+import os
+from collections import Counter
+
+from model import Model
+from utils import interact, Map
+
+class Optimizer(object):
+    def __init__(self, args, model):
+        self.args = args
+
+        self.save_dir = os.path.join(self.args.save_dir, 'optim')
+        os.makedirs(self.save_dir, exist_ok=True)
+
+        if isinstance(model, Model):
+            model = model.model
+
+        # set base arguments
+        kwargs_optimizer = {
+            'lr': args.lr,
+            'weight_decay': args.weight_decay
+        }
+
+        if args.optimizer == 'SGD':
+            optimizer_class = optim.SGD
+            kwargs_optimizer['momentum'] = args.momentum
+        elif args.optimizer == 'ADAM':
+            optimizer_class = optim.Adam
+            kwargs_optimizer['betas'] = args.betas
+            kwargs_optimizer['eps'] = args.epsilon
+        elif args.optimizer == 'RMSPROP':
+            optimizer_class = optim.RMSprop
+            kwargs_optimizer['eps'] = args.epsilon
+
+        # scheduler
+        if args.scheduler == 'step':
+            scheduler_class = lrs.MultiStepLR
+            kwargs_scheduler = {
+                'milestones': args.milestones,
+                'gamma': args.gamma,
+            }
+        elif args.scheduler == 'plateau':
+            scheduler_class = lrs.ReduceLROnPlateau
+            kwargs_scheduler = {
+                'mode': 'min',
+                'factor': args.gamma,
+                'patience': 10,
+                'verbose': True,
+                'threshold': 0,
+                'threshold_mode': 'abs',
+                'cooldown': 10,
+            }
+
+        self.kwargs_optimizer = kwargs_optimizer
+        self.scheduler_class = scheduler_class
+        self.kwargs_scheduler = kwargs_scheduler
+
+        def _get_optimizer(model):
+
+            class _Optimizer(optimizer_class):
+                def __init__(self, model, args, scheduler_class, kwargs_scheduler):
+                    trainable = filter(lambda x: x.requires_grad, model.parameters())
+                    super(_Optimizer, self).__init__(trainable, **kwargs_optimizer)
+
+                    self.args = args
+
+                    self._register_scheduler(scheduler_class, kwargs_scheduler)
+
+                def _register_scheduler(self, scheduler_class, kwargs_scheduler):
+                    self.scheduler = scheduler_class(self, **kwargs_scheduler)
+
+                def schedule(self, metrics=None):
+                    if isinstance(self, lrs.ReduceLROnPlateau):
+                        self.scheduler.step(metrics)
+                    else:
+                        self.scheduler.step()
+
+                def get_last_epoch(self):
+                    return self.scheduler.last_epoch
+
+                def get_lr(self):
+                    return self.param_groups[0]['lr']
+
+                def get_last_lr(self):
+                    return self.scheduler.get_last_lr()[0]
+
+                def state_dict(self):
+                    state_dict = super(_Optimizer, self).state_dict()  # {'state': ..., 'param_groups': ...}
+                    state_dict['scheduler'] = self.scheduler.state_dict()
+
+                    return state_dict
+
+                def load_state_dict(self, state_dict, epoch=None):
+                    # optimizer
+                    super(_Optimizer, self).load_state_dict(state_dict)  # load 'state' and 'param_groups' only
+                    # scheduler
+                    self.scheduler.load_state_dict(state_dict['scheduler']) # should work for plateau or simple resuming
+
+                    reschedule = False
+                    if isinstance(self.scheduler, lrs.MultiStepLR):
+                        if self.args.milestones != list(self.scheduler.milestones) or self.args.gamma != self.scheduler.gamma:
+                            reschedule = True
+
+                    if reschedule:
+                        if epoch is None:
+                            if self.scheduler.last_epoch > 1:
+                                epoch = self.scheduler.last_epoch
+                            else:
+                                epoch = self.args.start_epoch - 1
+
+                        # if False:
+                        #     # option 1. new scheduler
+                        #     for i, group in enumerate(self.param_groups):
+                        #         self.param_groups[i]['lr'] = group['initial_lr']    # reset optimizer learning rate to initial
+                        #     # self.scheduler = None
+                        #     self._register_scheduler(scheduler_class, kwargs_scheduler)
+
+                        #     self.zero_grad()
+                        #     self.step()
+                        #     for _ in range(epoch):
+                        #         self.scheduler.step()
+                        #     self._step_count -= 1
+
+                        # else:
+                        # option 2. modify existing scheduler
+                        self.scheduler.milestones = Counter(self.args.milestones)
+                        self.scheduler.gamma = self.args.gamma
+                        for i, group in enumerate(self.param_groups):
+                            self.param_groups[i]['lr'] = group['initial_lr']    # reset optimizer learning rate to initial
+                            multiplier = 1
+                            for milestone in self.scheduler.milestones:
+                                if epoch >= milestone:
+                                    multiplier *= self.scheduler.gamma
+
+                            self.param_groups[i]['lr'] *= multiplier
+
+            return _Optimizer(model, args, scheduler_class, kwargs_scheduler)
+
+        self.G = _get_optimizer(model.G)
+        if model.D is not None:
+            self.D = _get_optimizer(model.D)
+        else:
+            self.D = None
+
+        self.load(args.load_epoch)
+
+    def zero_grad(self):
+        self.G.zero_grad()
+
+    def step(self):
+        self.G.step()
+
+    def schedule(self, metrics=None):
+        self.G.schedule(metrics)
+        if self.D is not None:
+            self.D.schedule(metrics)
+
+    def get_last_epoch(self):
+        return self.G.get_last_epoch()
+
+    def get_lr(self):
+        return self.G.get_lr()
+
+    def get_last_lr(self):
+        return self.G.get_last_lr()
+
+    def state_dict(self):
+        state_dict = Map()
+        state_dict.G = self.G.state_dict()
+        if self.D is not None:
+            state_dict.D = self.D.state_dict()
+
+        return state_dict.toDict()
+
+    def load_state_dict(self, state_dict, epoch=None):
+        state_dict = Map(**state_dict)
+        self.G.load_state_dict(state_dict.G, epoch)
+        if self.D is not None:
+            self.D.load_state_dict(state_dict.D, epoch)
+
+    def _save_path(self, epoch=None):
+        epoch = epoch if epoch is not None else self.get_last_epoch()
+        save_path = os.path.join(self.save_dir, 'optim-{:d}.pt'.format(epoch))
+
+        return save_path
+
+    def save(self, epoch=None):
+        if epoch is None:
+            epoch = self.G.scheduler.last_epoch
+        torch.save(self.state_dict(), self._save_path(epoch))
+
+    def load(self, epoch):
+        if epoch > 0:
+            print('Loading optimizer from {}'.format(self._save_path(epoch)))
+            self.load_state_dict(torch.load(self._save_path(epoch), map_location=self.args.device), epoch=epoch)
+
+        elif epoch == 0:
+            pass
+        else:
+            raise NotImplementedError
+
+        return
+

deblur/src/optim/warm_multi_step_lr.py

@@ -0,0 +1,32 @@
+import math
+from bisect import bisect_right
+from torch.optim.lr_scheduler import _LRScheduler
+
+# MultiStep learning rate scheduler with warm restart
+class WarmMultiStepLR(_LRScheduler):
+    def __init__(self, optimizer, milestones, gamma=0.1, last_epoch=-1, scale=1):
+        if not list(milestones) == sorted(milestones):
+            raise ValueError(
+                'Milestones should be a list of increasing integers. Got {}',
+                milestones
+            )
+
+        self.milestones = milestones
+        self.gamma = gamma
+        self.scale = scale
+
+        self.warmup_epochs = 5
+        self.gradual = (self.scale - 1) / self.warmup_epochs
+        super(WarmMultiStepLR, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if self.last_epoch < self.warmup_epochs:
+            return [
+                base_lr * (1 + self.last_epoch * self.gradual) / self.scale
+                for base_lr in self.base_lrs
+            ]
+        else:
+            return [
+                base_lr * self.gamma ** bisect_right(self.milestones, self.last_epoch)
+                for base_lr in self.base_lrs
+            ]

deblur/src/option.py

@@ -0,0 +1,278 @@
+"""optionional argument parsing"""
+# pylint: disable=C0103, C0301
+import argparse
+import datetime
+import os
+import re
+import shutil
+import time
+
+import torch
+import torch.distributed as dist
+import torch.backends.cudnn as cudnn
+
+from utils import interact
+from utils import str2bool, int2str
+
+import template
+
+# Training settings
+parser = argparse.ArgumentParser(description='Dynamic Scene Deblurring')
+
+# Device specifications
+group_device = parser.add_argument_group('Device specs')
+group_device.add_argument('--seed', type=int, default=-1, help='random seed')
+group_device.add_argument('--num_workers', type=int, default=7, help='the number of dataloader workers')
+group_device.add_argument('--device_type', type=str, choices=('cpu', 'cuda'), default='cuda', help='device to run models')
+group_device.add_argument('--device_index', type=int, default=0, help='device id to run models')
+group_device.add_argument('--n_GPUs', type=int, default=1, help='the number of GPUs for training')
+group_device.add_argument('--distributed', type=str2bool, default=False, help='use DistributedDataParallel instead of DataParallel for better speed')
+group_device.add_argument('--launched', type=str2bool, default=False, help='identify if main.py was executed from launch.py. Do not set this to be true using main.py.')
+
+group_device.add_argument('--master_addr', type=str, default='127.0.0.1', help='master address for distributed')
+group_device.add_argument('--master_port', type=int2str, default='8023', help='master port for distributed')
+group_device.add_argument('--dist_backend', type=str, default='nccl', help='distributed backend')
+group_device.add_argument('--init_method', type=str, default='env://', help='distributed init method URL to discover peers')
+group_device.add_argument('--rank', type=int, default=0, help='rank of the distributed process (gpu id). 0 is the master process.')
+group_device.add_argument('--world_size', type=int, default=1, help='world_size for distributed training (number of GPUs)')
+
+# Data
+group_data = parser.add_argument_group('Data specs')
+group_data.add_argument('--data_root', type=str, default='/data/ssd/public/czli/deblur', help='dataset root location')
+group_data.add_argument('--dataset', type=str, default=None, help='training/validation/test dataset name, has priority if not None')
+group_data.add_argument('--data_train', type=str, default='GOPRO_Large', help='training dataset name')
+group_data.add_argument('--data_val', type=str, default=None, help='validation dataset name')
+group_data.add_argument('--data_test', type=str, default='GOPRO_Large', help='test dataset name')
+group_data.add_argument('--blur_key', type=str, default='blur_gamma', choices=('blur', 'blur_gamma'), help='blur type from camera response function for GOPRO_Large dataset')
+group_data.add_argument('--rgb_range', type=int, default=255, help='RGB pixel value ranging from 0')
+
+# Model
+group_model = parser.add_argument_group('Model specs')
+group_model.add_argument('--model', type=str, default='RecLamResNet', help='model architecture')
+group_model.add_argument('--pretrained', type=str, default='', help='pretrained model location')
+group_model.add_argument('--n_scales', type=int, default=5, help='multi-scale deblurring level')
+group_model.add_argument('--detach', type=str2bool, default=False, help='detach between recurrence')
+group_model.add_argument('--gaussian_pyramid', type=str2bool, default=True, help='gaussian pyramid input/target')
+group_model.add_argument('--n_resblocks', type=int, default=19, help='number of residual blocks per scale')
+group_model.add_argument('--n_feats', type=int, default=64, help='number of feature maps')
+group_model.add_argument('--kernel_size', type=int, default=5, help='size of conv kernel')
+group_model.add_argument('--downsample', type=str, choices=('Gaussian', 'bicubic', 'stride'), default='Gaussian', help='input pyramid generation method')
+
+group_model.add_argument('--precision', type=str, default='single', choices=('single', 'half'), help='FP precision for test(single | half)')
+
+# amp
+group_amp = parser.add_argument_group('AMP specs')
+group_amp.add_argument('--amp', type=str2bool, default=False, help='use automatic mixed precision training')
+group_amp.add_argument('--init_scale', type=float, default=1024., help='initial loss scale')
+
+# Training
+group_train = parser.add_argument_group('Training specs')
+group_train.add_argument('--patch_size', type=int, default=256, help='training patch size')
+group_train.add_argument('--batch_size', type=int, default=16, help='input batch size for training')
+group_train.add_argument('--split_batch', type=int, default=1, help='split a minibatch into smaller chunks')
+group_train.add_argument('--augment', type=str2bool, default=True, help='train with data augmentation')
+
+# Testing
+group_test = parser.add_argument_group('Testing specs')
+group_test.add_argument('--validate_every', type=int, default=10, help='do validation at every N epochs')
+group_test.add_argument('--test_every', type=int, default=10, help='do test at every N epochs')
+# group_test.add_argument('--chop', type=str2bool, default=False, help='memory-efficient forward')
+# group_test.add_argument('--self_ensemble', type=str2bool, default=False, help='self-ensembled testing')
+
+# Action
+group_action = parser.add_argument_group('Source behavior')
+group_action.add_argument('--do_train', type=str2bool, default=True, help='do train the model')
+group_action.add_argument('--do_validate', type=str2bool, default=True, help='do validate the model')
+group_action.add_argument('--do_test', type=str2bool, default=True, help='do test the model')
+group_action.add_argument('--demo', type=str2bool, default=False, help='demo')
+group_action.add_argument('--demo_input_dir', type=str, default='', help='demo input directory')
+group_action.add_argument('--demo_output_dir', type=str, default='', help='demo output directory')
+
+# Optimization
+group_optim = parser.add_argument_group('Optimization specs')
+group_optim.add_argument('--lr', type=float, default=1e-4, help='learning rate')
+group_optim.add_argument('--milestones', type=int, nargs='+', default=[500, 750, 900], help='learning rate decay per N epochs')
+group_optim.add_argument('--scheduler', default='step', choices=('step', 'plateau'), help='learning rate scheduler type')
+group_optim.add_argument('--gamma', type=float, default=0.5, help='learning rate decay factor for step decay')
+group_optim.add_argument('--optimizer', default='ADAM', choices=('SGD', 'ADAM', 'RMSprop'), help='optimizer to use (SGD | ADAM | RMSProp)')
+group_optim.add_argument('--momentum', type=float, default=0.9, help='SGD momentum')
+group_optim.add_argument('--betas', type=float, nargs=2, default=(0.9, 0.999), help='ADAM betas')
+group_optim.add_argument('--epsilon', type=float, default=1e-8, help='ADAM epsilon')
+group_optim.add_argument('--weight_decay', type=float, default=0, help='weight decay')
+group_optim.add_argument('--clip', type=float, default=0, help='weight decay')
+
+# Loss
+group_loss = parser.add_argument_group('Loss specs')
+group_loss.add_argument('--loss', type=str, default='1*MSE', help='loss function configuration')
+group_loss.add_argument('--metric', type=str, default='PSNR,SSIM', help='metric function configuration. ex) None | PSNR | SSIM | PSNR,SSIM')
+group_loss.add_argument('--gamma', type=float, default=0.6, help='gamma decay')
+
+
+# Logging
+group_log = parser.add_argument_group('Logging specs')
+group_log.add_argument('--save_dir', type=str, default='', help='subdirectory to save experiment logs')
+# group_log.add_argument('--load_dir', type=str, default='', help='subdirectory to load experiment logs')
+group_log.add_argument('--start_epoch', type=int, default=-1, help='(re)starting epoch number')
+group_log.add_argument('--end_epoch', type=int, default=1000, help='ending epoch number')
+group_log.add_argument('--load_epoch', type=int, default=-1, help='epoch number to load model (start_epoch-1 for training, start_epoch for testing)')
+group_log.add_argument('--save_every', type=int, default=10, help='save model/optimizer at every N epochs')
+group_log.add_argument('--save_results', type=str, default='part', choices=('none', 'part', 'all'), help='save none/part/all of result images')
+
+# Debugging
+group_debug = parser.add_argument_group('Debug specs')
+group_debug.add_argument('--stay', type=str2bool, default=False, help='stay at interactive console after trainer initialization')
+
+parser.add_argument('--template', type=str, default='', help='argument template option')
+
+args = parser.parse_args()
+template.set_template(args)
+
+args.data_root = os.path.expanduser(args.data_root)   # recognize home directory
+now = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+if args.save_dir == '':
+    args.save_dir = now
+args.save_dir = os.path.join('../experiment', args.save_dir)
+os.makedirs(args.save_dir, exist_ok=True)
+
+if args.start_epoch < 0: # start from scratch or continue from the last epoch
+    # check if there are any models saved before
+    model_dir = os.path.join(args.save_dir, 'models')
+    model_prefix = 'model-'
+    if os.path.exists(model_dir):
+        model_list = [name for name in os.listdir(model_dir) if name.startswith(model_prefix)]
+        last_epoch = 0
+        for name in model_list:
+            epochNumber = int(re.findall('\\d+', name)[0]) # model example name model-100.pt
+            if last_epoch < epochNumber:
+                last_epoch = epochNumber
+
+        args.start_epoch = last_epoch + 1
+    else:
+        # train from scratch
+        args.start_epoch = 1
+elif args.start_epoch == 0:
+    # remove existing directory and start over
+    if args.rank == 0:  # maybe local rank
+        shutil.rmtree(args.save_dir, ignore_errors=True)
+    os.makedirs(args.save_dir, exist_ok=True)
+    args.start_epoch = 1
+
+if args.load_epoch < 0:  # load_epoch == start_epoch when doing a post-training test for a specific epoch
+    args.load_epoch = args.start_epoch - 1
+
+if args.pretrained:
+    if args.start_epoch <= 1:
+        args.pretrained = os.path.join('../experiment', args.pretrained)
+    else:
+        print('starting from epoch {}! ignoring pretrained model path..'.format(args.start_epoch))
+        args.pretrained = ''
+
+if args.model == 'MSResNet':
+    args.gaussian_pyramid = True
+
+argname = os.path.join(args.save_dir, 'args.pt')
+argname_txt = os.path.join(args.save_dir, 'args.txt')
+if args.start_epoch > 1:
+    # load previous arguments and keep the necessary ones same
+
+    if os.path.exists(argname):
+        args_old = torch.load(argname)
+
+        load_list = []  # list of arguments that are fixed
+        # training
+        load_list += ['patch_size']
+        load_list += ['batch_size']
+        # data format
+        load_list += ['rgb_range']
+        load_list += ['blur_key']
+        # model architecture
+        load_list += ['n_scales']
+        load_list += ['n_resblocks']
+        load_list += ['n_feats']
+
+        for arg_part in load_list:
+            vars(args)[arg_part] = vars(args_old)[arg_part]
+
+if args.dataset is not None:
+    args.data_train = args.dataset
+    args.data_val = args.dataset if args.dataset != 'GOPRO_Large' else None
+    args.data_test = args.dataset
+
+if args.data_val is None:
+    args.do_validate = False
+
+if args.demo_input_dir:
+    args.demo = True
+
+if args.demo:
+    assert os.path.basename(args.save_dir) != now, 'You should specify pretrained directory by setting --save_dir SAVE_DIR'
+
+    args.data_train = ''
+    args.data_val = ''
+    args.data_test = ''
+
+    args.do_train = False
+    args.do_validate = False
+    args.do_test = False
+
+    assert len(args.demo_input_dir) > 0, 'Please specify demo_input_dir!'
+    args.demo_input_dir = os.path.expanduser(args.demo_input_dir)
+    if args.demo_output_dir:
+        args.demo_output_dir = os.path.expanduser(args.demo_output_dir)
+
+    args.save_results = 'all'
+
+if args.amp:
+    args.precision = 'single'   # model parameters should stay in fp32
+
+if args.seed < 0:
+    args.seed = int(time.time())
+
+# save arguments
+if args.rank == 0:
+    torch.save(args, argname)
+    with open(argname_txt, 'a') as file:
+        file.write('execution at {}\n'.format(now))
+
+        for key in args.__dict__:
+            file.write(key + ': ' + str(args.__dict__[key]) + '\n')
+
+        file.write('\n')
+
+# device and type
+if args.device_type == 'cuda' and not torch.cuda.is_available():
+    raise Exception("GPU not available!")
+
+if not args.distributed:
+    args.rank = 0
+
+def setup(args):
+    cudnn.benchmark = True
+
+    if args.distributed:
+        os.environ['MASTER_ADDR'] = args.master_addr
+        os.environ['MASTER_PORT'] = args.master_port
+
+        args.device_index = args.rank
+        args.world_size = args.n_GPUs   # consider single-node training
+
+        # initialize the process group
+        dist.init_process_group(args.dist_backend, init_method=args.init_method, rank=args.rank, world_size=args.world_size)
+
+    args.device = torch.device(args.device_type, args.device_index)
+    args.dtype = torch.float32
+    args.dtype_eval = torch.float32 if args.precision == 'single' else torch.float16
+
+    # set seed for processes (distributed: different seed for each process)
+    # model parameters are synchronized explicitly at initial
+    torch.manual_seed(args.seed)
+    if args.device_type == 'cuda':
+        torch.cuda.set_device(args.device)
+        if args.rank == 0:
+            torch.cuda.manual_seed_all(args.seed)
+
+    return args
+
+def cleanup(args):
+    if args.distributed:
+        dist.destroy_process_group()

deblur/src/prepare.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+workplace=`pwd -P`
+
+if [ ! -d "/data/ssd/public/czli/deblur/GOPRO_Large" ]; then
+    workplace=`pwd`
+    echo "Copying dataset to ssd"
+    cd /research/dept7/liuhy/deblur/dataset
+    mkdir -p /data/ssd/public/czli/deblur/GOPRO_Large
+    cp GOPRO_Large.zip /data/ssd/public/czli/deblur/GOPRO_Large
+    cd /data/ssd/public/czli/deblur/GOPRO_Large
+    echo "Dumping zip in data path:" `pwd`
+    for f in *.zip; do unzip "$f"; done
+fi
+
+cd "$workplace"
+echo "Workplace:" `pwd`

deblur/src/template.py

@@ -0,0 +1,9 @@
+def set_template(args):
+    if args.template.find('gopro') >= 0:
+        args.dataset = 'GOPRO_Large'
+        args.milestones = [500, 750, 900]
+        args.end_epoch = 1000
+    elif args.template.find('reds') >= 0:
+        args.dataset = 'REDS'
+        args.milestones = [100, 150, 180]
+        args.end_epoch = 200

deblur/src/train-rec-1.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+bash ./prepare.sh
+python launch.py --n_GPUs 2 main.py --batch_size 16 --amp true --save_dir LamRes_L1 --n_scales 1 --loss 1*L1

deblur/src/train.py

@@ -0,0 +1,249 @@
+import os
+from tqdm import tqdm
+
+import torch
+
+import data.common
+from utils import interact, MultiSaver
+from torch.utils.tensorboard import SummaryWriter
+import torchvision
+
+import torch.cuda.amp as amp
+
+class Trainer():
+
+    def __init__(self, args, model, criterion, optimizer, loaders):
+        print('===> Initializing trainer')
+        self.args = args
+        self.mode = 'train' # 'val', 'test'
+        self.epoch = args.start_epoch
+        self.save_dir = args.save_dir
+
+        self.model = model
+        self.criterion = criterion
+        self.optimizer = optimizer
+        self.loaders = loaders
+
+
+        self.do_train = args.do_train
+        self.do_validate = args.do_validate
+        self.do_test = args.do_test
+
+        self.device = args.device
+        self.dtype = args.dtype
+        self.dtype_eval = torch.float32 if args.precision == 'single' else torch.float16
+        self.recurrence=args.n_scales
+        if self.args.demo and self.args.demo_output_dir:
+            self.result_dir = self.args.demo_output_dir
+        else:
+            self.result_dir = os.path.join(self.save_dir, 'result')
+        os.makedirs(self.result_dir, exist_ok=True)
+        print('results are saved in {}'.format(self.result_dir))
+
+        self.imsaver = MultiSaver(self.result_dir)
+
+        self.is_slave = self.args.launched and self.args.rank != 0
+
+        self.scaler = amp.GradScaler(
+            init_scale=self.args.init_scale,
+            enabled=self.args.amp
+        )
+        if not self.is_slave:
+            self.writter=SummaryWriter(f"runs/{args.save_dir}")
+
+
+    def save(self, epoch=None):
+        epoch = self.epoch if epoch is None else epoch
+        if epoch % self.args.save_every == 0:
+            if self.mode == 'train':
+                self.model.save(epoch)
+                self.optimizer.save(epoch)
+            self.criterion.save()
+
+        return
+
+    def load(self, epoch=None, pretrained=None):
+        if epoch is None:
+            epoch = self.args.load_epoch
+        self.epoch = epoch
+        self.model.load(epoch, pretrained)
+        self.optimizer.load(epoch)
+        self.criterion.load(epoch)
+
+        return
+
+    def train(self, epoch):
+        self.mode = 'train'
+        self.epoch = epoch
+
+        self.model.train()
+        self.model.to(dtype=self.dtype)
+
+        self.criterion.train()
+        self.criterion.epoch = epoch
+
+        if not self.is_slave:
+            print('[Epoch {} / lr {:.2e}]'.format(
+                epoch, self.optimizer.get_lr()
+            ))
+        total=len(self.loaders[self.mode])
+        acc=0.0
+        if self.args.distributed:
+            self.loaders[self.mode].sampler.set_epoch(epoch)
+        if self.is_slave:
+            tq = self.loaders[self.mode]
+        else:
+            tq = tqdm(self.loaders[self.mode], ncols=80, smoothing=0, bar_format='{desc}|{bar}{r_bar}')
+        buffer=[0.0]*self.recurrence
+        torch.set_grad_enabled(True)
+        for idx, batch in enumerate(tq):
+            self.optimizer.zero_grad()
+
+            input, target = data.common.to(
+                batch[0], batch[1], device=self.device, dtype=self.dtype)
+
+            with amp.autocast(self.args.amp):
+                output = self.model(input)
+                loss = self.criterion(output, target)
+            
+            for i in range(self.recurrence):
+                buffer[i]+=self.criterion.buffer[i]
+            
+            self.scaler.scale(loss).backward()
+            if self.args.clip>0:
+                self.scaler.unscale_(self.optimizer.G)
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip)
+            self.scaler.step(self.optimizer.G)
+            self.scaler.update()
+
+            if isinstance(tq, tqdm):
+                tq.set_description(self.criterion.get_loss_desc())
+        if not self.is_slave:
+            rgb_range=self.args.rgb_range
+            for i in range(len(output)):
+                grid=torchvision.utils.make_grid(output[i])
+                self.writter.add_image(f"Output{i}",grid/rgb_range,epoch)
+                self.writter.add_scalar(f"Loss{i}",buffer[i],epoch)
+            self.writter.add_image("Input",torchvision.utils.make_grid(input[0])/rgb_range,epoch)
+            self.writter.add_image("Target",torchvision.utils.make_grid(target[0])/rgb_range,epoch)
+        self.criterion.normalize()
+        if isinstance(tq, tqdm):
+            tq.set_description(self.criterion.get_loss_desc())
+            tq.display(pos=-1)  # overwrite with synchronized loss
+
+        self.criterion.step()
+        self.optimizer.schedule(self.criterion.get_last_loss())
+
+        if self.args.rank == 0:
+            self.save(epoch)
+
+        return
+
+    def evaluate(self, epoch, mode='val'):
+        self.mode = mode
+        self.epoch = epoch
+
+        self.model.eval()
+        self.model.to(dtype=self.dtype_eval)
+
+        if mode == 'val':
+            self.criterion.validate()
+        elif mode == 'test':
+            self.criterion.test()
+        self.criterion.epoch = epoch
+
+        self.imsaver.join_background()
+
+        if self.is_slave:
+            tq = self.loaders[self.mode]
+        else:
+            tq = tqdm(self.loaders[self.mode], ncols=80, smoothing=0, bar_format='{desc}|{bar}{r_bar}')
+
+        compute_loss = True
+        torch.set_grad_enabled(False)
+        for idx, batch in enumerate(tq):
+            input, target = data.common.to(
+                batch[0], batch[1], device=self.device, dtype=self.dtype_eval)
+            with amp.autocast(self.args.amp):
+                output = self.model(input)
+            
+            # if self.args.rgb_range==1:
+                # output=output*255
+                # target=target*255
+
+            if mode == 'demo':  # remove padded part
+                pad_width = batch[2]
+                output[0], _ = data.common.pad(output[0], pad_width=pad_width, negative=True)
+
+            if isinstance(batch[1], torch.BoolTensor):
+                compute_loss = False
+
+            if compute_loss:
+                self.criterion(output, target)
+                if isinstance(tq, tqdm):
+                    tq.set_description(self.criterion.get_loss_desc())
+
+            if self.args.save_results != 'none':
+                if isinstance(output, (list, tuple)):
+                    result = output[-1]  # select last output in a pyramid
+                elif isinstance(output, torch.Tensor):
+                    result = output
+
+                names = batch[-1]
+
+                if self.args.save_results == 'part' and compute_loss: # save all when GT not available
+                    indices = batch[-2]
+                    save_ids = [save_id for save_id, idx in enumerate(indices) if idx % 10 == 0]
+
+                    result = result[save_ids]
+                    names = [names[save_id] for save_id in save_ids]
+
+                self.imsaver.save_image(result, names)
+
+        if compute_loss:
+            self.criterion.normalize()
+            if isinstance(tq, tqdm):
+                tq.set_description(self.criterion.get_loss_desc())
+                tq.display(pos=-1)  # overwrite with synchronized loss
+
+            self.criterion.step()
+            if self.args.rank == 0:
+                self.save()
+
+        self.imsaver.end_background()
+
+    def validate(self, epoch):
+        self.evaluate(epoch, 'val')
+        return
+
+    def test(self, epoch):
+        self.evaluate(epoch, 'test')
+        return
+
+    def fill_evaluation(self, epoch, mode=None, force=False):
+        if epoch <= 0:
+            return
+
+        if mode is not None:
+            self.mode = mode
+
+        do_eval = force
+        if not force:
+            loss_missing = epoch not in self.criterion.loss_stat[self.mode]['Total']    # should it switch to all loss types?
+
+            metric_missing = False
+            for metric_type in self.criterion.metric:
+                if epoch not in self.criterion.metric_stat[mode][metric_type]:
+                    metric_missing = True
+
+            do_eval = loss_missing or metric_missing
+
+        if do_eval:
+            try:
+                self.load(epoch)
+                self.evaluate(epoch, self.mode)
+            except:
+                # print('saved model/optimizer at epoch {} not found!'.format(epoch))
+                pass
+
+        return

deblur/src/train.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+bash ./prepare.sh
+python launch.py --n_GPUs 2 main.py --batch_size 16 --amp true --save_dir RecLamRes_MSE

deblur/src/train01.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+bash ./prepare.sh
+CUDA_VISIBLE_DEVICES=0,1 python launch.py --n_GPUs 2 main.py --batch_size 16 --amp true --save_dir RecLamRes_MSE

deblur/src/train23.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+bash ./prepare.sh
+CUDA_VISIBLE_DEVICES=2,3 python launch.py --n_GPUs 2 main.py --batch_size 8 --amp true --save_dir RecLamRes_MSE_detach --detach true --master_port 8025

deblur/src/utils.py

@@ -0,0 +1,200 @@
+import readline
+import rlcompleter
+readline.parse_and_bind("tab: complete")
+import code
+import pdb
+
+import time
+import argparse
+import os
+import imageio
+import torch
+import torch.multiprocessing as mp
+
+# debugging tools
+def interact(local=None):
+    """interactive console with autocomplete function. Useful for debugging.
+    interact(locals())
+    """
+    if local is None:
+        local=dict(globals(), **locals())
+
+    readline.set_completer(rlcompleter.Completer(local).complete)
+    code.interact(local=local)
+
+def set_trace(local=None):
+    """debugging with pdb
+    """
+    if local is None:
+        local=dict(globals(), **locals())
+
+    pdb.Pdb.complete = rlcompleter.Completer(local).complete
+    pdb.set_trace()
+
+# timer
+class Timer():
+    """Brought from https://github.com/thstkdgus35/EDSR-PyTorch
+    """
+    def __init__(self):
+        self.acc = 0
+        self.tic()
+
+    def tic(self):
+        self.t0 = time.time()
+
+    def toc(self):
+        return time.time() - self.t0
+
+    def hold(self):
+        self.acc += self.toc()
+
+    def release(self):
+        ret = self.acc
+        self.acc = 0
+
+        return ret
+
+    def reset(self):
+        self.acc = 0
+
+
+# argument parser type casting functions
+def str2bool(val):
+    """enable default constant true arguments"""
+    # https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse
+    if isinstance(val, bool):
+        return val
+    elif val.lower() == 'true':
+        return True
+    elif val.lower() == 'false':
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected')
+
+def int2str(val):
+    """convert int to str for environment variable related arguments"""
+    if isinstance(val, int):
+        return str(val)
+    elif isinstance(val, str):
+        return val
+    else:
+        raise argparse.ArgumentTypeError('number value expected')
+
+
+# image saver using multiprocessing queue
+class MultiSaver():
+    def __init__(self, result_dir=None):
+        self.queue = None
+        self.process = None
+        self.result_dir = result_dir
+
+    def begin_background(self):
+        self.queue = mp.Queue()
+
+        def t(queue):
+            while True:
+                if queue.empty():
+                    continue
+                img, name = queue.get()
+                if name:
+                    try:
+                        basename, ext = os.path.splitext(name)
+                        if ext != '.png':
+                            name = '{}.png'.format(basename)
+                        imageio.imwrite(name, img)
+                    except Exception as e:
+                        print(e)
+                else:
+                    return
+
+        worker = lambda: mp.Process(target=t, args=(self.queue,), daemon=False)
+        cpu_count = min(8, mp.cpu_count() - 1)
+        self.process = [worker() for _ in range(cpu_count)]
+        for p in self.process:
+            p.start()
+
+    def end_background(self):
+        if self.queue is None:
+            return
+
+        for _ in self.process:
+            self.queue.put((None, None))
+
+    def join_background(self):
+        if self.queue is None:
+            return
+
+        while not self.queue.empty():
+            time.sleep(0.5)
+
+        for p in self.process:
+            p.join()
+
+        self.queue = None
+
+    def save_image(self, output, save_names, result_dir=None):
+        result_dir = result_dir if self.result_dir is None else self.result_dir
+        if result_dir is None:
+            raise Exception('no result dir specified!')
+
+        if self.queue is None:
+            try:
+                self.begin_background()
+            except Exception as e:
+                print(e)
+                return
+
+        # assume NCHW format
+        if output.ndim == 2:
+            output = output.expand([1, 1] + list(output.shape))
+        elif output.ndim == 3:
+            output = output.expand([1] + list(output.shape))
+
+        for output_img, save_name in zip(output, save_names):
+            # assume image range [0, 255]
+            output_img = output_img.add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
+
+            save_name = os.path.join(result_dir, save_name)
+            save_dir = os.path.dirname(save_name)
+            os.makedirs(save_dir, exist_ok=True)
+
+            self.queue.put((output_img, save_name))
+
+        return
+
+class Map(dict):
+    """
+    https://stackoverflow.com/questions/2352181/how-to-use-a-dot-to-access-members-of-dictionary
+    Example:
+    m = Map({'first_name': 'Eduardo'}, last_name='Pool', age=24, sports=['Soccer'])
+    """
+    def __init__(self, *args, **kwargs):
+        super(Map, self).__init__(*args, **kwargs)
+        for arg in args:
+            if isinstance(arg, dict):
+                for k, v in arg.items():
+                    self[k] = v
+
+        if kwargs:
+            for k, v in kwargs.items():
+                self[k] = v
+
+    def __getattr__(self, attr):
+        return self.get(attr)
+
+    def __setattr__(self, key, value):
+        self.__setitem__(key, value)
+
+    def __setitem__(self, key, value):
+        super(Map, self).__setitem__(key, value)
+        self.__dict__.update({key: value})
+
+    def __delattr__(self, item):
+        self.__delitem__(item)
+
+    def __delitem__(self, key):
+        super(Map, self).__delitem__(key)
+        del self.__dict__[key]
+
+    def toDict(self):
+        return self.__dict__