Upload 10 files

carn-pcsr-phase1.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfe84ddd3923b35d14a977dabda5613a9d59da3b0961e004be786f108d3f8508
+size 755333

demo.py

@@ -0,0 +1,58 @@
+import torch
+import models
+from torchvision import transforms
+from utils import *
+from PIL import Image
+import numpy as np
+
+img_path = '/workspace/datasets/test/myimage/HR/X4/FOTO-BOX-18-1024x1024.png' # only support .png
+scale = 4 # only support x4
+
+'''
+k: hyperparameter to traverse PSNR-FLOPs trade-off. smaller k → larger FLOPs & PSNR. range is about [-1,2].
+adaptive: whether to use automatic decision of k
+no_refinement: whether not to use pixel-wise refinement (postprocessing for reducing artifacts)
+parser.add_argument('--opacity', type=float, default=0.65, help='opacity for colored visualization')
+parser.add_argument('--pixel_batch_size', type=int, default=300000)
+'''
+
+resume_path = 'carn-pcsr-phase1.pth'
+sv_file = torch.load(resume_path)
+model = models.make(sv_file['model'], load_sd=True).cuda()
+model.eval()
+
+rgb_mean = torch.tensor([0.4488, 0.4371, 0.4040], device='cuda').view(1,3,1,1)
+rgb_std = torch.tensor([1.0, 1.0, 1.0], device='cuda').view(1,3,1,1)
+
+with torch.no_grad():
+    # prepare inputs
+    lr = transforms.ToTensor()(Image.open(img_path)).unsqueeze(0).cuda() # (1,3,h,w), range=[0,1]
+    h,w = lr.shape[-2:]
+    H,W = h*scale, w*scale
+    coord = make_coord((H,W), flatten=True, device='cuda').unsqueeze(0)
+    cell = torch.ones_like(coord)
+    cell[:,:,0] *= 2/H
+    cell[:,:,1] *= 2/W
+    inp_lr = (lr - rgb_mean) / rgb_std
+
+    pred, flag = model(inp_lr, coord=coord, cell=cell, scale=scale, k=0, 
+        pixel_batch_size=300000, adaptive_cluster=True, refinement=True)
+    flops = get_model_flops(model, inp_lr, coord=coord, cell=cell, scale=scale, k=0, 
+        pixel_batch_size=300000, adaptive_cluster=True, refinement=True)
+    max_flops = get_model_flops(model, inp_lr, coord=coord, cell=cell, scale=scale, k=-25, 
+        pixel_batch_size=300000, adaptive_cluster=False, refinement=True)
+    print('flops: {:.1f}G ({:.1f} %) | max_flops: {:.1f}G (100 %)'.format(flops/1e9, 
+        (flops / max_flops)*100, max_flops/1e9))
+
+pred = pred.transpose(1,2).view(-1,3,H,W)
+pred = pred * rgb_std + rgb_mean
+pred = tensor2numpy(pred)
+Image.fromarray(pred).save(f'output.png')
+
+flag = flag.view(-1,1,H,W).repeat(1,3,1,1).squeeze(0).detach().cpu()
+H,W = pred.shape[:2]
+vis_img = np.zeros_like(pred)
+vis_img[flag[0] == 0] = np.array([0,255,0])
+vis_img[flag[0] == 1] = np.array([255,0,0])
+vis_img = vis_img*0.35 + pred*0.65
+Image.fromarray(vis_img.astype('uint8')).save('output_vis.png')

models/__init__.py

@@ -0,0 +1,3 @@
+from .models import register, make
+from . import mlp, pcsr, sampler
+from . import carn

models/carn.py

@@ -0,0 +1,78 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import models.utils as mutils
+from models import register
+
+
+class Block(nn.Module):
+    def __init__(self, nf, group=1):
+        super(Block, self).__init__()
+        self.b1 = mutils.EResidualBlock(nf, nf, group=group)
+        self.c1 = mutils.BasicBlock(nf*2, nf, 1, 1, 0)
+        self.c2 = mutils.BasicBlock(nf*3, nf, 1, 1, 0)
+        self.c3 = mutils.BasicBlock(nf*4, nf, 1, 1, 0)
+
+    def forward(self, x):
+        c0 = o0 = x
+
+        b1 = self.b1(o0)
+        c1 = torch.cat([c0, b1], dim=1)
+        o1 = self.c1(c1)
+        
+        b2 = self.b1(o1)
+        c2 = torch.cat([c1, b2], dim=1)
+        o2 = self.c2(c2)
+        
+        b3 = self.b1(o2)
+        c3 = torch.cat([c2, b3], dim=1)
+        o3 = self.c3(c3)
+
+        return o3
+
+
+@register('carn')
+class CARN_M(nn.Module):
+    def __init__(self, in_nc=3, out_nc=3, nf=64, scale=4, group=4, no_upsampling=False):
+        super(CARN_M, self).__init__()
+        self.scale = scale
+        self.out_dim = nf
+
+        self.entry = nn.Conv2d(in_nc, nf, 3, 1, 1)
+        self.b1 = Block(nf, group=group)
+        self.b2 = Block(nf, group=group)
+        self.b3 = Block(nf, group=group)
+
+        self.c1 = mutils.BasicBlock(nf*2, nf, 1, 1, 0)
+        self.c2 = mutils.BasicBlock(nf*3, nf, 1, 1, 0)
+        self.c3 = mutils.BasicBlock(nf*4, nf, 1, 1, 0)
+
+        self.no_upsampling = no_upsampling
+        if not no_upsampling:
+            self.upsample = mutils.UpsampleBlock(nf, scale=scale, multi_scale=False, group=group)
+            self.exit = nn.Conv2d(nf, out_nc, 3, 1, 1)
+                
+    def forward(self, x):
+        #x = self.sub_mean(x)
+        x = self.entry(x)
+        c0 = o0 = x
+
+        b1 = self.b1(o0)
+        c1 = torch.cat([c0, b1], dim=1)
+        o1 = self.c1(c1)
+        
+        b2 = self.b2(o1)
+        c2 = torch.cat([c1, b2], dim=1)
+        o2 = self.c2(c2)
+        
+        b3 = self.b3(o2)
+        c3 = torch.cat([c2, b3], dim=1)
+        o3 = self.c3(c3)
+        out = o3.clone()
+
+        if not self.no_upsampling:
+            out = self.upsample(out, scale=self.scale)
+            out = self.exit(out)
+        #out = self.add_mean(out)
+        return out

models/mlp.py

@@ -0,0 +1,32 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from models import register
+
+
+@register('mlp')
+class MLP(nn.Module):
+
+    def __init__(self, in_dim, out_dim, hidden_list, residual=False):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.hidden_list = hidden_list
+        self.residual = residual
+        if residual:
+            self.convert = nn.Linear(in_dim, out_dim)
+
+        layers = []
+        lastv = in_dim
+        for hidden in hidden_list:
+            layers.append(nn.Linear(lastv, hidden))
+            layers.append(nn.ReLU())
+            lastv = hidden
+        layers.append(nn.Linear(lastv, out_dim))
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        y = self.layers(x)
+        if self.residual:
+            y = y + self.convert(x)
+        return y

models/models.py

@@ -0,0 +1,23 @@
+import copy
+
+
+models = {}
+
+
+def register(name):
+    def decorator(cls):
+        models[name] = cls
+        return cls
+    return decorator
+
+
+def make(model_spec, args=None, load_sd=False):
+    if args is not None:
+        model_args = copy.deepcopy(model_spec['args'])
+        model_args.update(args)
+    else:
+        model_args = model_spec['args']
+    model = models[model_spec['name']](**model_args)
+    if load_sd:
+        model.load_state_dict(model_spec['sd'])
+    return model

models/pcsr.py

@@ -0,0 +1,197 @@
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import models
+from models import register
+from fast_pytorch_kmeans import KMeans
+from utils import *
+
+
+@register('pcsr-phase0')
+class PCSR(nn.Module):
+    def __init__(self, encoder_spec, heavy_sampler_spec):
+        super().__init__()
+        self.encoder = models.make(encoder_spec)
+        in_dim = self.encoder.out_dim
+        self.heavy_sampler = models.make(heavy_sampler_spec,
+            args={'in_dim': in_dim, 'out_dim': 3})
+
+    def forward(self, lr, coord, cell, **kwargs):
+        if self.training:
+            return self.forward_train(lr, coord, cell)
+        else:
+            return self.forward_test(lr, coord, cell, **kwargs)
+
+    def forward_train(self, lr, coord, cell):
+        feat = self.encoder(lr)
+        res = F.grid_sample(lr, coord.flip(-1).unsqueeze(1), mode='bilinear', 
+            padding_mode='border', align_corners=False)[:,:,0,:].permute(0,2,1)
+        pred_heavy = self.heavy_sampler(feat, coord, cell) + res
+        return pred_heavy
+
+    def forward_test(self, lr, coord, cell, pixel_batch_size=None):
+        feat = self.encoder(lr)
+        b,q = coord.shape[:2]
+        tot = b*q
+        if not pixel_batch_size:
+            pixel_batch_size = q
+
+        preds = []
+        for i in range(b): # for each image
+            pred = torch.zeros((q,3), device=lr.device)
+            l = 0
+            while l < q:
+                r = min(q, l+pixel_batch_size)
+                coord_split = coord[i:i+1,l:r,:]
+                cell_split = cell[i:i+1,l:r,:]
+                res = F.grid_sample(lr[i:i+1], coord_split.flip(-1).unsqueeze(1), mode='bilinear', 
+                    padding_mode='border', align_corners=False)[:,:,0,:].squeeze(0).transpose(0,1)
+                pred[l:r] = self.heavy_sampler(feat[i:i+1], coord_split, cell_split) + res
+                l = r
+            preds.append(pred)
+        pred = torch.stack(preds, dim=0)
+        return pred
+
+
+@register('pcsr-phase1')
+class PCSR(nn.Module):
+
+    def __init__(self, encoder_spec, heavy_sampler_spec, light_sampler_spec, classifier_spec):
+        super().__init__()
+        self.encoder = models.make(encoder_spec)
+        in_dim = self.encoder.out_dim
+        self.heavy_sampler = models.make(heavy_sampler_spec,
+            args={'in_dim': in_dim, 'out_dim': 3})
+        self.light_sampler = models.make(light_sampler_spec,
+            args={'in_dim': in_dim, 'out_dim': 3})
+        self.classifier = models.make(classifier_spec,
+            args={'in_dim': in_dim, 'out_dim': 2})
+        self.kmeans = KMeans(n_clusters=2, max_iter=20, mode='euclidean', verbose=0)
+        self.cost_list = {}
+
+    def forward(self, lr, coord, cell, **kwargs):
+        if self.training:
+            return self.forward_train(lr, coord, cell)
+        else:
+            return self.forward_test(lr, coord, cell, **kwargs)
+
+    def forward_train(self, lr, coord, cell):
+        feat = self.encoder(lr)
+        prob = self.classifier(feat, coord, cell)
+        prob = F.softmax(prob, dim=-1) # (b,q,2)
+
+        pred_heavy = self.heavy_sampler(feat, coord, cell)
+        pred_light = self.light_sampler(feat, coord, cell)
+        pred = prob[:,:,0:1] * pred_light + prob[:,:,1:2] * pred_heavy
+
+        res = F.grid_sample(lr, coord.flip(-1).unsqueeze(1), mode='bilinear', 
+            padding_mode='border', align_corners=False)[:,:,0,:].permute(0,2,1)
+        pred = pred + res
+        return pred, prob
+
+    def forward_test(self, lr, coord, cell, scale=None, hr_size=None, k=0., pixel_batch_size=None, adaptive_cluster=False, refinement=True):
+        h,w = lr.shape[-2:]
+        if not scale and hr_size:
+            H,W = hr_size
+            scale = round((H/h + W/w)/2, 1)
+        else:
+            assert scale and not hr_size
+            H,W = round(h*scale), round(w*scale)
+            hr_size = (H,W)
+
+        if scale not in self.cost_list:
+            h0,w0 = 16,16
+            H0,W0 = round(h0*scale), round(w0*scale)
+            inp_coord = make_coord((H0,W0), flatten=True, device='cuda').unsqueeze(0)
+            inp_cell = torch.ones_like(inp_coord)
+            inp_cell[:,:,0] *= 2/H0
+            inp_cell[:,:,1] *= 2/W0
+            inp_encoder = torch.zeros((1,3,h0,w0), device='cuda')
+            flops_encoder = get_model_flops(self.encoder, inp_encoder)
+            inp_sampler = torch.zeros((1,self.encoder.out_dim,h0,w0), device='cuda')
+            x = get_model_flops(self.light_sampler, inp_sampler, coord=inp_coord, cell=inp_cell)
+            y = get_model_flops(self.heavy_sampler, inp_sampler, coord=inp_coord, cell=inp_cell)
+            cost_list = torch.FloatTensor([x,y]).cuda() + flops_encoder
+            cost_list = cost_list / cost_list.sum()
+            self.cost_list[scale] = cost_list
+            print('cost_list calculated (x{}): {}'.format(scale, cost_list))
+        cost_list = self.cost_list[scale]
+
+        feat = self.encoder(lr)
+        b,q = coord.shape[:2]
+        assert H*W == q
+        tot = b*q
+        if not pixel_batch_size: 
+            pixel_batch_size = q
+
+        # pre-calculate flag
+        prob = torch.zeros((b,q,2), device=lr.device)
+        pb = pixel_batch_size//b*b
+        assert pb > 0
+        l = 0
+        while l < q:
+            r = min(q, l+pb)
+            coord_split = coord[:,l:r,:]
+            cell_split = cell[:,l:r,:]
+            prob_split = self.classifier(feat, coord_split, cell_split)
+            prob[:,l:r] = F.softmax(prob_split, dim=-1)
+            l = r      
+
+        if adaptive_cluster: # auto-decide threshold
+            diff = prob[:,:,1].view(-1,1) # (tot,1)
+            assert diff.max() > diff.min()
+            diff = (diff - diff.min()) / (diff.max() - diff.min())
+            centroids = torch.FloatTensor([[0.5]]).cuda()
+            flag = self.kmeans.fit_predict(diff, centroids=centroids)
+            _, min_index = torch.min(diff.flatten(), dim=0)
+            if flag[min_index] == 1:
+                flag = 1 - flag # (tot,)
+            flag = flag.view(b,q)
+        else:
+            prob = prob / torch.pow(cost_list, k).view(1,1,2)
+            flag = torch.argmax(prob, dim=-1) # (b,q)
+
+        # inference per image
+        # more efficient implementation may exist
+        preds = []
+        for i in range(b):
+            pred = torch.zeros((q,3), device=lr.device)
+            l = 0
+            while l < q:
+                r = min(q, l+pixel_batch_size)
+                coord_split = coord[i:i+1,l:r,:]
+                cell_split = cell[i:i+1,l:r,:]
+                flg = flag[i,l:r]
+
+                idx_easy = torch.where(flg == 0)[0]
+                idx_hard = torch.where(flg == 1)[0]
+                num_easy, num_hard = len(idx_easy), len(idx_hard)
+                if num_easy > 0: 
+                    pred[l+idx_easy] = self.light_sampler(feat[i:i+1], coord_split[:,idx_easy,:], cell_split[:,idx_easy,:]).squeeze(0)
+                if num_hard > 0: 
+                    pred[l+idx_hard] = self.heavy_sampler(feat[i:i+1], coord_split[:,idx_hard,:], cell_split[:,idx_hard,:]).squeeze(0)
+                res = F.grid_sample(lr[i:i+1], coord_split.flip(-1).unsqueeze(1), mode='bilinear', 
+                    padding_mode='border', align_corners=False)[:,:,0,:].squeeze(0).transpose(0,1)
+                pred[l:r] += res
+                l = r
+            preds.append(pred)
+        pred = torch.stack(preds, dim=0) # (b,q,3)
+
+        if refinement:
+            pred = pred.transpose(1,2).view(-1,3,H,W)
+            pred_unfold = F.pad(pred, (1,1,1,1), mode='replicate')
+            pred_unfold = F.unfold(pred_unfold, 3, padding=0).view(-1,3,9,H,W).mean(dim=2) # (b,3,H,W)
+            flag = flag.view(-1,1,H,W)
+            flag_unfold = F.pad(flag.float(), (1,1,1,1), mode='replicate')
+            flag_unfold = F.unfold(flag_unfold, 3, padding=0).view(-1,1,9,H,W).int().sum(dim=2) # (b,1,H,W)
+
+            cond = (flag==0) & (flag_unfold>0) #
+            cond[:,:,[0,-1],:] = cond[:,:,:,[0,-1]] = False
+            #print('refined: {} / {}'.format(cond.sum().item(), tot))
+            pred = torch.where(cond, pred_unfold, pred)
+            pred = pred.view(-1,3,q).transpose(1,2)
+        flag = flag.view(b,q,1)
+        return pred, flag

models/sampler.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import models
+from models import register
+from utils import make_coord
+
+
+@register('liif-sampler')
+class LIIF_Sampler(nn.Module):
+    # feature unfolding, local ensemble not supported
+    def __init__(self, imnet_spec, in_dim, out_dim):
+        super().__init__()
+        self.imnet = models.make(imnet_spec, args={'in_dim': in_dim+4, 'out_dim': out_dim})
+
+    def make_inp(self, feat, coord, cell):
+        feat_coord = make_coord(feat.shape[-2:], flatten=False, device=feat.device)\
+            .permute(2,0,1).unsqueeze(0).expand(feat.shape[0], 2, *feat.shape[-2:])
+        q_feat = F.grid_sample(feat, coord.flip(-1).unsqueeze(1), mode='nearest',
+            align_corners=False)[:,:,0,:].permute(0,2,1)
+        q_coord = F.grid_sample(feat_coord, coord.flip(-1).unsqueeze(1), mode='nearest',
+            align_corners=False)[:,:,0,:].permute(0,2,1)
+
+        rel_coord = coord - q_coord
+        rel_coord[:,:,0] *= feat.shape[-2]
+        rel_coord[:,:,1] *= feat.shape[-1]
+
+        rel_cell = cell.clone()
+        rel_cell[:,:,0] *= feat.shape[-2]
+        rel_cell[:,:,1] *= feat.shape[-1]
+
+        inp = torch.cat([q_feat, rel_coord, rel_cell], dim=-1)
+        return inp    
+
+    def forward(self, x, coord=None, cell=None):
+        if coord is not None:
+            x = self.make_inp(x, coord, cell)
+        x = self.imnet(x)
+        return x

models/utils.py

@@ -0,0 +1,214 @@
+# https://github.com/XPixelGroup/ClassSR
+import math
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+import torch.nn.functional as F
+
+
+def initialize_weights(net_l, scale=1):
+    if not isinstance(net_l, list):
+        net_l = [net_l]
+    for net in net_l:
+        for m in net.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, a=0, mode='fan_in')
+                m.weight.data *= scale  # for residual block
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                init.kaiming_normal_(m.weight, a=0, mode='fan_in')
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias.data, 0.0)
+
+
+def make_layer(block, n_layers):
+    layers = []
+    for _ in range(n_layers):
+        layers.append(block())
+    return nn.Sequential(*layers)
+
+
+class ResidualBlock_noBN(nn.Module):
+    '''Residual block w/o BN
+    ---Conv-ReLU-Conv-+-
+     |________________|
+    '''
+
+    def __init__(self, nf=64):
+        super(ResidualBlock_noBN, self).__init__()
+        self.conv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+
+        # initialization
+        initialize_weights([self.conv1, self.conv2], 0.1)
+
+    def forward(self, x):
+        identity = x
+        out = F.relu(self.conv1(x), inplace=True)
+        out = self.conv2(out)
+        return identity + out
+
+
+def default_conv(in_channels, out_channels, kernel_size, bias=True):
+    return nn.Conv2d(
+        in_channels, out_channels, kernel_size,
+        padding=(kernel_size//2), bias=bias)
+
+class MeanShift(nn.Conv2d):
+    def __init__(self, rgb_range, rgb_mean, rgb_std, sign=-1):
+        super(MeanShift, self).__init__(3, 3, kernel_size=1)
+        std = torch.Tensor(rgb_std)
+        self.weight.data = torch.eye(3).view(3, 3, 1, 1)
+        self.weight.data.div_(std.view(3, 1, 1, 1))
+        self.bias.data = sign * rgb_range * torch.Tensor(rgb_mean)
+        self.bias.data.div_(std)
+        self.requires_grad = False
+
+class BasicBlock(nn.Sequential):
+    def __init__(
+        self, in_channels, out_channels, kernel_size, stride=1, bias=False,
+        bn=True, act=nn.ReLU(True)):
+
+        m = [nn.Conv2d(
+            in_channels, out_channels, kernel_size,
+            padding=(kernel_size//2), stride=stride, bias=bias)
+        ]
+        if bn: m.append(nn.BatchNorm2d(out_channels))
+        if act is not None: m.append(act)
+        super(BasicBlock, self).__init__(*m)
+
+class ResBlock(nn.Module):
+    def __init__(
+        self, conv, n_feat, kernel_size,
+        bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
+
+        super(ResBlock, self).__init__()
+        m = []
+        for i in range(2):
+            m.append(conv(n_feat, n_feat, kernel_size, bias=bias))
+            if bn: m.append(nn.BatchNorm2d(n_feat))
+            if i == 0: m.append(act)
+
+        self.body = nn.Sequential(*m)
+        self.res_scale = res_scale
+
+    def forward(self, x):
+        res = self.body(x).mul(self.res_scale)
+        res += x
+
+        return res
+
+class Upsampler(nn.Sequential):
+    def __init__(self, conv, scale, n_feat, bn=False, act=False, bias=True):
+
+        m = []
+        if (scale & (scale - 1)) == 0:    # Is scale = 2^n?
+            for _ in range(int(math.log(scale, 2))):
+                m.append(conv(n_feat, 4 * n_feat, 3, bias))
+                m.append(nn.PixelShuffle(2))
+                if bn: m.append(nn.BatchNorm2d(n_feat))
+                if act: m.append(act())
+        elif scale == 3:
+            m.append(conv(n_feat, 9 * n_feat, 3, bias))
+            m.append(nn.PixelShuffle(3))
+            if bn: m.append(nn.BatchNorm2d(n_feat))
+            if act: m.append(act())
+        else:
+            raise NotImplementedError
+
+        super(Upsampler, self).__init__(*m)
+
+
+class EResidualBlock(nn.Module):
+    def __init__(self,
+                 in_channels, out_channels,
+                 group=1):
+        super(EResidualBlock, self).__init__()
+
+        self.body = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, 1, 1, groups=group),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, 3, 1, 1, groups=group),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, 1, 1, 0),
+        )
+
+    def forward(self, x):
+        out = self.body(x)
+        out = F.relu(out + x)
+        return out
+
+
+class Upsampler(nn.Sequential):
+    def __init__(self, conv, scale, n_feat, bn=False, act=False, bias=True):
+
+        m = []
+        if (scale & (scale - 1)) == 0:  # Is scale = 2^n?
+            for _ in range(int(math.log(scale, 2))):
+                m.append(conv(n_feat, 4 * n_feat, 3, bias))
+                m.append(nn.PixelShuffle(2))
+                if bn: m.append(nn.BatchNorm2d(n_feat))
+                if act: m.append(act())
+        elif scale == 3:
+            m.append(conv(n_feat, 9 * n_feat, 3, bias))
+            m.append(nn.PixelShuffle(3))
+            if bn: m.append(nn.BatchNorm2d(n_feat))
+            if act: m.append(act())
+        else:
+            raise NotImplementedError
+
+        super(Upsampler, self).__init__(*m)
+
+
+class UpsampleBlock(nn.Module):
+    def __init__(self,
+                 n_channels, scale, multi_scale,
+                 group=1):
+        super(UpsampleBlock, self).__init__()
+
+        if multi_scale:
+            self.up2 = _UpsampleBlock(n_channels, scale=2, group=group)
+            self.up3 = _UpsampleBlock(n_channels, scale=3, group=group)
+            self.up4 = _UpsampleBlock(n_channels, scale=4, group=group)
+        else:
+            self.up = _UpsampleBlock(n_channels, scale=scale, group=group)
+
+        self.multi_scale = multi_scale
+
+    def forward(self, x, scale):
+        if self.multi_scale:
+            if scale == 2:
+                return self.up2(x)
+            elif scale == 3:
+                return self.up3(x)
+            elif scale == 4:
+                return self.up4(x)
+        else:
+            return self.up(x)
+
+
+class _UpsampleBlock(nn.Module):
+    def __init__(self,
+                 n_channels, scale,
+                 group=1):
+        super(_UpsampleBlock, self).__init__()
+
+        modules = []
+        if scale == 2 or scale == 4 or scale == 8:
+            for _ in range(int(math.log(scale, 2))):
+                modules += [nn.Conv2d(n_channels, 4 * n_channels, 3, 1, 1, groups=group), nn.ReLU(inplace=True)]
+                modules += [nn.PixelShuffle(2)]
+        elif scale == 3:
+            modules += [nn.Conv2d(n_channels, 9 * n_channels, 3, 1, 1, groups=group), nn.ReLU(inplace=True)]
+            modules += [nn.PixelShuffle(3)]
+
+        self.body = nn.Sequential(*modules)
+
+    def forward(self, x):
+        out = self.body(x)
+        return out

utils.py

@@ -0,0 +1,210 @@
+import numpy as np
+import torch
+import torch.nn as nn
+
+from collections import OrderedDict
+import pandas as pd
+import warnings
+warnings.filterwarnings("ignore")
+
+def tensor2numpy(tensor, rgb_range=1.):
+    rgb_coefficient = 255 / rgb_range
+    img = tensor.mul(rgb_coefficient).clamp(0, 255).round()
+    img = img[0].data if img.ndim==4 else img.data
+    img = np.transpose(img.cpu().numpy(), (1, 2, 0)).astype(np.uint8)
+    return img
+
+def center_crop(img, size):
+    h,w = img.shape[-2:]
+    cut_h, cut_w = h-size[0], w-size[1]
+    
+    lh = cut_h // 2
+    rh = h - (cut_h - lh)
+    lw = cut_w // 2
+    rw = w - (cut_w - lw)
+    
+    img = img[:,:, lh:rh, lw:rw]
+    return img
+
+def make_coord(shape, ranges=None, flatten=True, device='cpu'):
+    # Make coordinates at grid centers.
+    coord_seqs = []
+    for i, n in enumerate(shape):
+        if ranges is None:
+            v0, v1 = -1, 1
+        else:
+            v0, v1 = ranges[i]
+        r = (v1 - v0) / (2 * n)
+        seq = v0 + r + (2 * r) * torch.arange(n, device=device).float()
+        coord_seqs.append(seq)
+    ret = torch.stack(torch.meshgrid(*coord_seqs), dim=-1)
+    if flatten:
+        ret = ret.view(-1, ret.shape[-1])
+    return ret
+
+def compute_num_params(model, text=False):
+    tot = int(sum([np.prod(p.shape) for p in model.parameters()]))
+    if text:
+        if tot >= 1e6:
+            return '{:.3f}M'.format(tot / 1e6)
+        elif tot >= 1e3:
+            return '{:.2f}K'.format(tot / 1e3)
+        else:
+            return '{}'.format(tot)
+    else:
+        return tot
+    
+
+def get_names_dict(model):
+    """Recursive walk to get names including path."""
+    names = {}
+
+    def _get_names(module, parent_name=""):
+        for key, m in module.named_children():
+            cls_name = str(m.__class__).split(".")[-1].split("'")[0]
+            num_named_children = len(list(m.named_children()))
+            if num_named_children > 0:
+                name = parent_name + "." + key if parent_name else key
+            else:
+                name = parent_name + "." + cls_name + "_"+ key if parent_name else key
+            names[name] = m
+
+            if isinstance(m, nn.Module):
+                _get_names(m, parent_name=name)
+
+    _get_names(model)
+    return names
+
+# https://github.com/chenbong/ARM-Net/blob/main/utils/util.py
+def get_model_flops(model, x, *args, **kwargs):
+    """Summarize the given input model.
+    Summarized information are 1) output shape, 2) kernel shape,
+    3) number of the parameters and 4) operations (Mult-Adds)
+    Args:
+        model (Module): Model to summarize
+        x (Tensor): Input tensor of the model with [N, C, H, W] shape
+                    dtype and device have to match to the model
+        args, kwargs: Other argument used in `model.forward` function
+    """
+    model.eval()
+    if hasattr(model, 'module'):
+        model = model.module
+    #x = torch.zeros(input_size).to(next(model.parameters()).device)
+
+    def register_hook(module):
+        def hook(module, inputs, outputs):
+            cls_name = str(module.__class__).split(".")[-1].split("'")[0]
+            module_idx = len(summary)
+            key = None
+            for name, item in module_names.items():
+                if item == module:
+                    key = "{}_{}".format(module_idx, name)
+                    break
+            assert key
+
+            info = OrderedDict()
+            info["id"] = id(module)
+            if isinstance(outputs, (list, tuple)):
+                try:
+                    info["out"] = list(outputs[0].size())
+                except AttributeError:
+                    info["out"] = list(outputs[0].data.size())
+            else:
+                info["out"] = list(outputs.size())
+
+            info["ksize"] = "-"
+            info["inner"] = OrderedDict()
+            info["params_nt"], info["params"], info["flops"] = 0, 0, 0
+
+            for name, param in module.named_parameters():
+                info["params"] += param.nelement() * param.requires_grad
+                info["params_nt"] += param.nelement() * (not param.requires_grad)
+
+                if name == "weight":
+                    ksize = list(param.size())
+                    if len(ksize) > 1:
+                        ksize[0], ksize[1] = ksize[1], ksize[0]
+                    info["ksize"] = ksize
+
+                    if isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d):
+                        assert len(inputs[0].size()) == 4 and len(inputs[0].size()) == len(outputs[0].size())+1
+
+                        in_c, in_h, in_w = inputs[0].size()[1:]
+                        k_h, k_w = module.kernel_size
+                        out_c, out_h, out_w = outputs[0].size()
+                        groups = module.groups
+                        kernel_mul = k_h * k_w * (in_c // groups)
+
+                        kernel_mul_group = kernel_mul * out_h * out_w * (out_c // groups)
+                        total_mul = kernel_mul_group * groups
+                        info["flops"] += 2 * total_mul * inputs[0].size()[0] # total
+
+                    elif isinstance(module, nn.BatchNorm2d):
+                        info["flops"] += 2 * inputs[0].numel()
+
+                    elif isinstance(module, nn.InstanceNorm2d):
+                        info["flops"] += 6 * inputs[0].numel()
+
+                    elif isinstance(module, nn.LayerNorm):
+                        info["flops"] += 8 * inputs[0].numel()
+
+                    elif isinstance(module, nn.Linear):
+                        q = inputs[0].numel() // inputs[0].shape[-1]
+                        info["flops"] += 2*q * module.in_features * module.out_features # total
+
+                    elif isinstance(module, nn.PReLU) or isinstance(module, nn.ReLU):
+                        info["flops"] += inputs[0].numel()
+                    else:
+                        print('not supported:', module)
+                        exit()
+                        info["flops"] += param.nelement()
+
+                elif "weight" in name:
+                    info["inner"][name] = list(param.size())
+                    info["flops"] += param.nelement()
+
+            if list(module.named_parameters()):
+                for v in summary.values():
+                    if info["id"] == v["id"]:
+                        info["params"] = "(recursive)"
+
+            #if info["params"] == 0:
+            #    info["params"], info["flops"] = "-", "-"
+            summary[key] = info
+
+        if not module._modules:
+            hooks.append(module.register_forward_hook(hook))
+
+    module_names = get_names_dict(model)
+    hooks = []
+    summary = OrderedDict()
+
+    model.apply(register_hook)
+    try:
+        with torch.no_grad():
+            model(x) if not (kwargs or args) else model(x, *args, **kwargs)
+    finally:
+        for hook in hooks:
+            hook.remove()
+    # Use pandas to align the columns
+    df = pd.DataFrame(summary).T
+
+    df["Mult-Adds"] = pd.to_numeric(df["flops"], errors="coerce")
+    df["Params"] = pd.to_numeric(df["params"], errors="coerce")
+    df["Non-trainable params"] = pd.to_numeric(df["params_nt"], errors="coerce")
+    df = df.rename(columns=dict(
+        ksize="Kernel Shape",
+        out="Output Shape",
+    ))
+    return df['Mult-Adds'].sum()
+    '''
+    with warnings.catch_warnings():
+        warnings.filterwarnings('ignore')
+        df_sum = df.sum()
+
+    df.index.name = "Layer"
+
+    df = df[["Kernel Shape", "Output Shape", "Params", "Mult-Adds"]]
+    max_repr_width = max([len(row) for row in df.to_string().split("\n")])
+    return df_sum["Mult-Adds"]
+    '''