spikezoo 0.1.2__py3-none-any.whl → 0.2.1__py3-none-any.whl

spikezoo/archs/bsf/main.py

@@ -0,0 +1,398 @@
+import argparse
+import os
+import os.path as osp
+import time
+import numpy as np
+import torch
+import torch.backends.cudnn as cudnn
+from tensorboardX import SummaryWriter
+import datetime
+from datasets import datasets
+from models.get_model import get_model
+from utils import *
+from metrics.psnr import *
+from metrics.ssim import *
+import lpips
+from skimage.metrics import peak_signal_noise_ratio as compare_psnr
+import pprint
+from models.bsf.dsft_convert import convert_dsft4
+
+parser = argparse.ArgumentParser()
+############################ Dataset Root ############################
+parser.add_argument('--dataset_storage', type=str, default='ram')   ## ram or disk
+parser.add_argument('--data-root', type=str, default='/dev/shm/rzhao/REDS120fps')
+parser.add_argument('--half_reserve', type=int, default=2, help=' DSFT half reserve + 3ref + 4key + 3ref + DSFT half reserve')
+############################ Training Params ############################
+parser.add_argument('--arch', '-a', type=str, default='MEPF')
+parser.add_argument('--batch-size', '-bs', type=int, default=8)
+parser.add_argument('--learning-rate', '-lr', type=float, default=2e-4)
+parser.add_argument('--train-res', '-tr', type=int, default=[128, 128], metavar='N', nargs='*')
+parser.add_argument('--input-type', type=str, default='raw_spike', choices=['dsft', 'raw_spike'])
+parser.add_argument('--epochs', '-ep', type=int, default=100)
+parser.add_argument('--workers', '-j', type=int, default=8)
+parser.add_argument('--pretrained', '-prt', type=str, default=None)
+parser.add_argument('--start-epoch', '-sep', type=int, default=0)
+parser.add_argument('--print-freq', '-pf', type=int, default=100)
+parser.add_argument('--save-dir', '-sd', type=str, default='outputs')
+parser.add_argument('--save-name', '-sn', type=str, default=None)
+parser.add_argument('--vis-path', '-vp', type=str, default='vis')
+parser.add_argument('--vis-name', '-vn', type=str, default='model1')
+parser.add_argument('--eval-path', '-evp', type=str, default='eval_vis/model1')
+parser.add_argument('--vis-freq', '-vf', type=int, default=20)
+parser.add_argument('--eval', '-e', action='store_true')
+parser.add_argument('--print_details', '-pd', action='store_true')
+parser.add_argument('--milestones', default=[10, 20, 30, 40, 50, 60, 70, 80, 90, 100], metavar='N', nargs='*')
+parser.add_argument('--lr-scale-factor', '-lrsf', type=float, default=0.7)
+parser.add_argument('--eval-interval', '-ei', type=int, default=5)
+parser.add_argument('--no_imwrite', action='store_true', default=False)
+parser.add_argument('--compile_model', '-cmpmd', action='store_true')
+parser.add_argument('--seed', type=int, default=2728)
+############################ Params about Dataset ############################
+parser.add_argument('--alpha', type=float, default=0.7)
+parser.add_argument('--eta_list', default=[1.00, 0.75, 0.50], type=float, metavar='N', nargs='*')
+parser.add_argument('--gamma', type=int, default=60)
+############################ About Optimizer ############################
+parser.add_argument('--solver', type=str, default='Adam')
+parser.add_argument('--momentum', type=float, default=0.9)
+parser.add_argument('--beta', type=float, default=0.999)
+parser.add_argument('--weight_decay', type=float, default=0.0)
+
+parser.add_argument('--test_eval', action='store_true')
+parser.add_argument('--logs_file_name', type=str, default='bsf')
+
+parser.add_argument('--loss_type', type=str, default='l1')
+parser.add_argument('--dsft_convertor_type', type=int, default=4)
+
+parser.add_argument('--no_dsft', action='store_true')
+args = parser.parse_args()
+
+
+
+##########################################################################################################
+## configs
+writer_root = 'logs/{:s}/'.format(args.logs_file_name)
+os.makedirs(writer_root, exist_ok=True)
+writer_path = writer_root + args.arch + '.txt'
+writer = open(writer_path, 'a')
+
+for k, v in vars(args).items():
+    vv = pprint.pformat(v)
+    ostr = '{:s} : {:s}'.format(k, vv)
+    writer.write(ostr + '\n')
+
+args.milestones = [int(m) for m in args.milestones]
+ostr = 'milsones '
+for mmm in args.milestones:
+    ostr += '{:d} '.format(mmm)
+writer.write(ostr + '\n')
+
+n_iter = 0
+
+
+def batch_PSNR(img, imclean, data_range):
+    Img = img.data.cpu().numpy().astype(np.float32)
+    Iclean = imclean.data.cpu().numpy().astype(np.float32)
+    PSNR = compare_psnr(Img, Iclean, data_range=data_range)
+    '''
+    PSNR = 0
+    for i in range(Img.shape[0]):
+        PSNR += compare_psnr(Iclean[i,:,:,:], Img[i,:,:,:], data_range=data_range)
+    return (PSNR/Img.shape[0])
+    '''
+    return PSNR
+
+
+def train(args, train_loader, model, optimizer, epoch, train_writer):
+    ######################################################################
+    ## Init
+    global n_iter
+    batch_time = AverageMeter(precision=3)
+    data_time = AverageMeter(precision=3)
+    
+    losses = AverageMeter(precision=6, names=['Loss'])
+    batch_psnr = AverageMeter(precision=4)
+
+    model.train()
+    
+    end = time.time()
+    
+    ######################################################################
+    ## Training Loop
+    for ww, data in enumerate(train_loader, 0):
+        spikes = [spk.float().cuda() for spk in data['spikes']]
+        spks = torch.cat(spikes, dim=1)
+        central_idx = 10*args.half_reserve + 30
+        spks = spks[:, central_idx-30:central_idx+31]
+
+        if not args.no_dsft:
+            dsfts = [d.float().cuda() for d in data['dsft']]
+            dsfts = torch.cat(dsfts, dim=1)
+            central_idx = 10*args.half_reserve + 30
+            dsfts = dsfts[:, central_idx-30:central_idx+31]
+
+        images = [img.cuda() for img in data['images']]
+        norm_fac = data['norm_fac'].unsqueeze_(dim=1).unsqueeze_(dim=1).unsqueeze_(dim=1).cuda().float()
+        
+        data_time.update(time.time() - end)
+
+        if not args.no_dsft:
+            dsft_dict = convert_dsft4(spike=spks, dsft=dsfts)
+
+            input_dict = {
+                'dsft_dict': dsft_dict,
+                'spikes': spks,
+            }
+
+        gt = images[0]
+
+        rec = model(input_dict=input_dict)
+
+        rec = rec / norm_fac
+
+        if args.loss_type == 'l1':
+            loss = (rec - gt).abs().mean()
+        elif args.loss_type == 'charbonnier':
+            loss = torch.sqrt((rec - gt)**2 + 1e-6).mean()
+
+        # record loss
+        losses.update(loss)
+        cur_batch_psnr = batch_PSNR(img=rec, imclean=gt, data_range=1.0)
+        batch_psnr.update(cur_batch_psnr)
+
+        if ww % 10 == 0:
+            train_writer.add_scalar('loss', loss.item(), n_iter)
+            train_writer.add_scalar('batch_psnr', cur_batch_psnr, n_iter)
+
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        
+        batch_time.update(time.time() - end)
+        n_iter += 1
+
+        if n_iter % args.vis_freq == 0:
+            vis_img(args.vis_path, torch.clip(rec, 0, 1), args.arch)
+        
+        ostr = 'Epoch: [{:03d}] [{:04d}/{:04d}],  Iter: {:6d}  '.format(epoch+1, ww, len(train_loader), n_iter-1)
+        ostr += 'Time: {},  Data: {}  '.format(batch_time, data_time)
+        ostr += ' '.join(map('{:s} {:.4f} ({:.6f}) '.format, losses.names, losses.val, losses.avg))
+        ostr += 'batch_PSNR {} '.format(batch_psnr)
+        ostr += 'lr {:.6f}'.format(optimizer.state_dict()['param_groups'][0]['lr'])
+        if ww % args.print_freq == 0:
+            writer.write(ostr + '\n')
+        end = time.time()
+    
+    return
+
+
+def validation(args, test_loader_list, model, lpips_function_dict):
+    model.eval()
+    
+    for eta, test_loader in zip(args.eta_list, test_loader_list):
+        cur_eval_root = osp.join(args.eval_path, args.arch, 'eta_{:.2f}'.format(eta))
+        os.makedirs(cur_eval_root, exist_ok=True)
+
+        global n_iter
+        batch_time = AverageMeter()
+        data_time = AverageMeter()
+        metrics_name = ['PSNR', 'SSIM', 'LPIPS-A', 'LPIPS-V', 'AvgTime']
+        metrics = AverageMeter(i=len(metrics_name), precision=4, names=metrics_name)
+        
+        for ww, data in enumerate(test_loader, 0):
+            st1 = time.time()
+            spks = torch.cat([spk.float().cuda() for spk in data['spikes']], dim=1)
+            central_idx = 10*args.half_reserve + 30
+            spks = spks[:, central_idx-30:central_idx+31]
+
+            if not args.no_dsft:
+                dsfts = torch.cat([d.float().cuda() for d in data['dsft']], dim=1)
+                central_idx = 10*args.half_reserve + 30
+                dsfts = dsfts[:, central_idx-30:central_idx+31]
+
+            images = data['images']
+            norm_fac = data['norm_fac'].unsqueeze_(dim=1).unsqueeze_(dim=1).unsqueeze_(dim=1).cuda().float()
+            
+            data_time.update(time.time() - st1)
+
+            if not args.no_dsft:
+                dsft_dict = convert_dsft4(spike=spks, dsft=dsfts)
+                    
+                input_dict = {
+                    'dsft_dict': dsft_dict,
+                    'spikes': spks,
+                }
+
+            with torch.no_grad():
+                st = time.time()
+                rec = model(input_dict=input_dict)
+                mtime = time.time() - st
+
+            rec = rec / norm_fac
+            rec = torch.clip(rec, 0, 1)
+            rec_np = torch2numpy255(rec)
+            img_np = torch2numpy255(images[0])
+
+            if not args.no_imwrite:
+                cur_vis_path = osp.join(cur_eval_root, '{:03d}.png'.format(ww))
+                cv2.imwrite(cur_vis_path, rec_np.astype(np.uint8))
+
+            cur_psnr = calculate_psnr(rec_np, img_np)
+            cur_ssim = calculate_ssim(rec_np, img_np)
+            with torch.no_grad():
+                cur_lpips_alex = lpips_function_dict['alex'](rec, images[0].cuda())
+                cur_lpips_vgg  = lpips_function_dict['vgg'](rec, images[0].cuda())
+
+            cur_metrics_list = [cur_psnr, cur_ssim, cur_lpips_alex.item(), cur_lpips_vgg.item() , mtime]
+            metrics.update(cur_metrics_list)
+
+        torch.cuda.empty_cache()
+        ostr = 'Eta {:.2f}  ALL  '.format(eta) + '  '.join(map('{:s} {:.4f}'.format, metrics.names, metrics.avg))
+        writer.write(ostr + '\n')
+
+    return
+    
+
+def main():
+    ##########################################################################################################
+    # Set random seeds
+    set_seeds(args.seed)
+
+    # Create save path and logs
+    timestamp1 = datetime.datetime.now().strftime('%m-%d')
+    timestamp2 = datetime.datetime.now().strftime('%H%M%S')
+
+    save_folder_name = 'a_{:s}_b{:d}_{:s}'.format(args.arch, args.batch_size, timestamp2)
+
+    save_path = osp.join(args.save_dir, timestamp1, save_folder_name)
+    make_dir(save_path)
+    ostr = '=>Save path: ' + save_path
+    writer.write(ostr + '\n')
+    # print('=>Save path: ', save_path)
+    train_writer = SummaryWriter(save_path)
+    
+    make_dir(args.vis_path)
+    make_dir(args.eval_path)    
+
+    model = None
+    optimizer = None
+
+    ##########################################################################################################
+    ## Create model
+    print(args.arch)
+    model = get_model(args)
+
+    if args.compile_model and (torch.__version__ >= '2.0.0'):
+        ostr = 'Start compile the model'
+        writer.write(ostr + '\n')
+        st = time.time()
+        torch.compile(model)
+        ostr = 'Finish compiling the model  Time {:.2f}s'.format(time.time() - st)
+        writer.write(ostr + '\n')
+
+    if args.pretrained != None:
+        network_data = torch.load(args.pretrained)
+        ostr = '=> using pretrained model {:s}'.format(args.pretrained)
+        writer.write(ostr + '\n')
+        ostr = '=> model params: {:.6f}M'.format(model.num_parameters()/1e6)
+        writer.write(ostr + '\n')
+        model = torch.nn.DataParallel(model).cuda()
+        model = model.cuda()
+        model.load_state_dict(network_data)
+    else:
+        network_data = None
+        ostr = '=> train from scratch'
+        writer.write(ostr + '\n')
+        model.init_weights()
+        ostr = '=> model params: {:.6f}M'.format(model.num_parameters()/1e6)
+        writer.write(ostr + '\n')
+        model = torch.nn.DataParallel(model).cuda()
+        model = model.cuda()
+
+    cudnn.benchmark = True
+
+    ##########################################################################################################
+    ## Create Optimizer
+    assert(args.solver in ['Adam', 'SGD'])
+    ostr = '=> settings {:s} solver'.format(args.solver)
+    writer.write(ostr + '\n')
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
+
+
+    ##########################################################################################################
+    ## Dataset
+    train_set = datasets.sreds_train(args)
+    train_loader = torch.utils.data.DataLoader(
+        train_set,
+        drop_last=False,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.workers,
+    )
+
+    test_loader_list = []
+    for eta in args.eta_list:
+        test_set = datasets.sreds_test_small(args, eta=eta)
+        test_loader = torch.utils.data.DataLoader(
+            test_set,
+            drop_last=False,
+            batch_size=1,
+            shuffle=False,
+            num_workers=args.workers,
+        )
+        test_loader_list.append(test_loader)
+
+    ##########################################################################################################
+    ## For LPIPS
+    loss_fn_alex = lpips.LPIPS(net='alex').cuda() # closer to "traditional" perceptual loss, when used for optimization
+    loss_fn_vgg = lpips.LPIPS(net='vgg').cuda() # closer to "traditional" perceptual loss, when used for optimization
+    lpips_function_dict = {'alex': loss_fn_alex, 'vgg': loss_fn_vgg}
+
+    ##########################################################################################################
+    ## Train or Evaluate
+    if args.test_eval:
+        validation(
+            args=args, 
+            test_loader_list=test_loader_list, 
+            model=model,
+            lpips_function_dict=lpips_function_dict,
+            )
+        return
+
+    epoch = args.start_epoch
+    while(True):
+        train(
+            args=args,
+            train_loader=train_loader,
+            model=model,
+            optimizer=optimizer,
+            epoch=epoch,
+            train_writer=train_writer,
+        )
+        epoch += 1
+
+        # scheduler can be added here
+        if epoch in args.milestones:
+            for param_group in optimizer.param_groups:
+                param_group['lr'] = param_group['lr'] * args.lr_scale_factor
+
+        # save model
+        if epoch % 5 == 0:
+            model_save_name = '{:s}_epoch{:03d}.pth'.format(args.arch, epoch)
+            torch.save(model.state_dict(), osp.join(save_path, model_save_name))
+        
+        # if epoch % 5 == 0:
+        if epoch % args.eval_interval == 0:
+            validation(
+                args=args, 
+                test_loader_list=test_loader_list, 
+                model=model,
+                lpips_function_dict=lpips_function_dict,
+                )
+
+
+        if epoch >= args.epochs:
+            break
+
+
+if __name__ == '__main__':
+    main()

spikezoo/archs/bsf/metrics/psnr.py

@@ -0,0 +1,22 @@
+import math
+import numpy as np
+
+# --------------------------------------------
+# PSNR
+# --------------------------------------------
+def calculate_psnr(img1, img2, border=0):
+    # img1 and img2 have range [0, 255]
+    #img1 = img1.squeeze()
+    #img2 = img2.squeeze()
+    if not img1.shape == img2.shape:
+        raise ValueError('Input images must have the same dimensions.')
+    h, w = img1.shape[:2]
+    img1 = img1[border:h-border, border:w-border]
+    img2 = img2[border:h-border, border:w-border]
+
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    mse = np.mean((img1 - img2)**2)
+    if mse == 0:
+        return float('inf')
+    return 20 * math.log10(255.0 / math.sqrt(mse))

spikezoo/archs/bsf/metrics/ssim.py

@@ -0,0 +1,54 @@
+import numpy as np
+import cv2
+
+# --------------------------------------------
+# SSIM
+# --------------------------------------------
+def calculate_ssim(img1, img2, border=0):
+    '''calculate SSIM
+    the same outputs as MATLAB's
+    img1, img2: [0, 255]
+    '''
+    #img1 = img1.squeeze()
+    #img2 = img2.squeeze()
+    if not img1.shape == img2.shape:
+        raise ValueError('Input images must have the same dimensions.')
+    h, w = img1.shape[:2]
+    img1 = img1[border:h-border, border:w-border]
+    img2 = img2[border:h-border, border:w-border]
+
+    if img1.ndim == 2:
+        return ssim(img1, img2)
+    elif img1.ndim == 3:
+        if img1.shape[2] == 3:
+            ssims = []
+            for i in range(3):
+                ssims.append(ssim(img1[:,:,i], img2[:,:,i]))
+            return np.array(ssims).mean()
+        elif img1.shape[2] == 1:
+            return ssim(np.squeeze(img1), np.squeeze(img2))
+    else:
+        raise ValueError('Wrong input image dimensions.')
+
+
+def ssim(img1, img2):
+    C1 = (0.01 * 255)**2
+    C2 = (0.03 * 255)**2
+
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    kernel = cv2.getGaussianKernel(11, 1.5)
+    window = np.outer(kernel, kernel.transpose())
+
+    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
+    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
+    mu1_sq = mu1**2
+    mu2_sq = mu2**2
+    mu1_mu2 = mu1 * mu2
+    sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
+    sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
+    sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
+
+    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
+                                                            (sigma1_sq + sigma2_sq + C2))
+    return ssim_map.mean()

spikezoo/archs/bsf/models/bsf/align.py

@@ -0,0 +1,154 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.ops import DeformConv2d
+
+
+class CAPA(nn.Module):
+    def __init__(self, chnn, sc=11):
+        super().__init__()
+        self.sc = sc
+        self.unfold = nn.Unfold(kernel_size=3*self.sc, dilation=1, padding=self.sc, stride=self.sc)
+        self.scale = chnn ** -0.5
+        self.to_q = nn.Conv2d(chnn, chnn, 1, bias=False)
+        self.to_k = nn.Conv2d(chnn, chnn, 1, bias=False)
+        self.to_v = nn.Conv2d(chnn, chnn, 1, bias=False)
+        self.gamma = nn.Parameter(torch.zeros(1))
+        self.mask_k = True
+
+    def forward(self, x_key, x_ref):
+        b, c, h_in, w_in = x_key.shape
+        x_pad = self.sc - w_in % self.sc
+        y_pad = self.sc - h_in % self.sc
+        feat_key = F.pad(x_key, (0, x_pad, 0, y_pad)) 
+        feat_ref = F.pad(x_ref, (0, x_pad, 0, y_pad)) 
+        b, c, h, w = feat_key.shape
+        h_sc = h // self.sc
+        w_sc = w // self.sc 
+
+        fm = torch.ones(1, 1, h_in, w_in).to(feat_key.device)
+        fm = F.pad(fm, (0, x_pad, 0, y_pad))
+        fm_k = self.unfold(fm).view(1, 1, -1, h_sc*w_sc)
+        fm_q = fm.view(1, 1, h_sc, self.sc, w_sc, self.sc).permute(0, 1, 2, 4, 3, 5).contiguous().view(1, 1, h_sc*w_sc, self.sc**2)
+        am = torch.einsum('b c k n, b c n s -> b k n s', fm_k, fm_q)
+        am = (am - 1) * 99.
+        am = am.repeat(b, 1, 1, 1)
+
+        feat_q = self.to_q(feat_key)
+        feat_k = self.to_k(feat_ref)
+        feat_k = self.unfold(feat_k).view(b, c, -1, h_sc*w_sc)
+        feat_k = self.scale * feat_k
+        feat_q = feat_q.view(b, c, h_sc, self.sc, w_sc, self.sc).permute(0, 1, 2, 4, 3, 5).contiguous().view(b, c, h_sc*w_sc, self.sc**2)
+        attn = torch.einsum('b c k n, b c n s -> b k n s', feat_k, feat_q)
+
+
+        attn = attn + am
+        self.attn = F.softmax(attn, dim=1)
+
+        feat_v = self.to_v(feat_ref)
+        feat_v = self.unfold(feat_v).view(b, c, -1, h_sc*w_sc)
+        feat_r = torch.einsum('b k n s, b c k n -> b c n s', self.attn, feat_v)
+        feat_r = feat_r.view(b, c, h_sc, w_sc, self.sc, self.sc).permute(0, 1, 2, 4, 3, 5).contiguous().view(b, c, h, w)
+        feat_r = feat_r[:,:,:h_in,:w_in]
+        feat_o = x_ref + feat_r * self.gamma
+        return feat_o
+
+
+class Multi_Granularity_Align_One_Level(nn.Module):
+    def __init__(self, base_dim=64, offset_groups=4, act=nn.ReLU(), memory=True):
+        super().__init__()
+        self.offset_groups = offset_groups
+        self.memory = memory
+
+        if self.memory:
+            first_output_dim = base_dim
+        else:
+            first_output_dim = 3*self.offset_groups*3*3
+
+        self.offset_conv_1 = self._make_two_conv_layer(input_dim=base_dim*2, hidden_dim=base_dim, output_dim=first_output_dim, kernel_size=3, stride=1, padding=1, act=act)
+        if self.memory:
+            self.offset_conv2_1 = self._make_two_conv_layer(input_dim=base_dim + 3*self.offset_groups*3*3, hidden_dim=base_dim, output_dim=3*self.offset_groups*3*3, kernel_size=3, stride=1, padding=1, act=act)
+            self.fuse_feat = self._make_two_conv_layer(input_dim=base_dim*2, hidden_dim=base_dim, output_dim=base_dim, kernel_size=3, stride=1, padding=1, act=act)
+
+    def _make_two_conv_layer(self, input_dim, hidden_dim, output_dim, kernel_size, stride, padding, act):
+        layer = nn.Sequential(
+            nn.Conv2d(input_dim, hidden_dim, kernel_size=kernel_size, stride=stride, padding=padding),
+            act,
+            nn.Conv2d(hidden_dim, output_dim, kernel_size=kernel_size, stride=1, padding=padding),
+        )
+        return layer
+
+    def forward(self, feat_list, kpa, deform_conv, prev_offset_feat=None, prev_feat=None):
+        xa, xb, xc, xd, xe = feat_list
+
+        xa_kpa, xb_kpa, xd_kpa, xe_kpa = [kpa(x_key=xc, x_ref=xxx) for xxx in [xa, xb, xd, xe]]
+        feat_for_conv_offset1 = [torch.cat([xxx, xc], dim=1) for xxx in [xa_kpa, xb_kpa, xd_kpa, xe_kpa]]
+
+        offset_feat_list1 = [self.offset_conv_1(f) for f in feat_for_conv_offset1]
+        if self.memory:
+            prev_offset_upsample_list1 = [F.interpolate(offset_feat, scale_factor=2, mode='bilinear') for offset_feat in prev_offset_feat]
+            offset_feat_list1 = [self.offset_conv2_1(torch.cat((f1, f2), dim=1)) for f1, f2 in zip(offset_feat_list1, prev_offset_upsample_list1)]
+        
+        o1o2m_abde_list1 = [f.chunk(3, dim=1) for f in offset_feat_list1]
+        offset_abde_list1 = [torch.cat((o1o2m[0], o1o2m[1]), dim=1) for o1o2m in o1o2m_abde_list1]
+        mask_abde_list1 = [torch.sigmoid(o1o2m[2]) for o1o2m in o1o2m_abde_list1]
+        
+        x_align_abde = [deform_conv(input=xx, offset=offset, mask=mask) for xx,offset,mask in zip([xa_kpa, xb_kpa, xd_kpa, xe_kpa], offset_abde_list1, mask_abde_list1)]
+
+        if self.memory:
+            prev_x_abde_align_upasmple_list = [F.interpolate(xxx, scale_factor=2, mode='bilinear') for xxx in prev_feat]
+            x_align_abde = [self.fuse_feat(torch.cat((x_align, prev_x_align_upsample), dim=1)) for x_align, prev_x_align_upsample in zip(x_align_abde, prev_x_abde_align_upasmple_list)]
+
+        xa_align, xb_align, xd_align, xe_align = x_align_abde
+        x_align = (xa_align, xb_align, xc, xd_align, xe_align)
+        return x_align, offset_feat_list1
+
+
+class Multi_Granularity_Align(nn.Module):
+    def __init__(self, base_dim=64, groups=4, act=nn.ReLU(), sc=11):
+        super().__init__()
+        self.offset_groups = groups
+        self.sc = sc
+
+        self.kpa = CAPA(chnn=base_dim, sc=self.sc)
+        self.deform_conv = DeformConv2d(in_channels=base_dim, out_channels=base_dim, kernel_size=3, stride=1, padding=1, groups=1)
+
+        ## Downsample
+        self.conv_ds_L2 = self._make_two_conv_layer(input_dim=base_dim, hidden_dim=base_dim, output_dim=base_dim, kernel_size=3, stride=2, padding=1, act=act)
+        self.conv_ds_L3 = self._make_two_conv_layer(input_dim=base_dim, hidden_dim=base_dim, output_dim=base_dim, kernel_size=3, stride=2, padding=1, act=act)
+
+        self.align_L3 = Multi_Granularity_Align_One_Level(base_dim=base_dim, offset_groups=self.offset_groups, act=act, memory=False)
+        self.align_L2 = Multi_Granularity_Align_One_Level(base_dim=base_dim, offset_groups=self.offset_groups, act=act, memory=True)
+        self.align_L1 = Multi_Granularity_Align_One_Level(base_dim=base_dim, offset_groups=self.offset_groups, act=act, memory=True)
+
+    def _make_two_conv_layer(self, input_dim, hidden_dim, output_dim, kernel_size, stride, padding, act):
+        layer = nn.Sequential(
+            nn.Conv2d(input_dim, hidden_dim, kernel_size=kernel_size, stride=stride, padding=padding),
+            act,
+            nn.Conv2d(hidden_dim, output_dim, kernel_size=kernel_size, stride=1, padding=padding),
+        )
+        return layer
+
+    def forward(self, feat_list):
+        '''
+        feat_list: xa, xb, xc, xd, xe
+        '''
+
+        xa, xb, xc, xd, xe = feat_list
+        ## Downsample
+        feat_batch_cat_L1 = torch.cat(feat_list, dim=0)
+        feat_batch_cat_L2 = self.conv_ds_L2(feat_batch_cat_L1)
+        feat_batch_cat_L3 = self.conv_ds_L3(feat_batch_cat_L2)
+
+        L3_align_feat_list, L3_offset_feat_list = self.align_L3(feat_list=feat_batch_cat_L3.chunk(5, dim=0), kpa=self.kpa, deform_conv=self.deform_conv,
+                                                                prev_offset_feat=None, prev_feat=None)
+
+        L2_align_feat_list, L2_offset_feat_list = self.align_L2(feat_list=feat_batch_cat_L2.chunk(5, dim=0), kpa=self.kpa, deform_conv=self.deform_conv, 
+                                                                prev_offset_feat=L3_offset_feat_list, prev_feat=L3_align_feat_list)
+
+        L1_align_feat_list, L1_offset_feat_list = self.align_L1(feat_list=feat_batch_cat_L1.chunk(5, dim=0), kpa=self.kpa, deform_conv=self.deform_conv,
+                                                                prev_offset_feat=L2_offset_feat_list, prev_feat=L2_align_feat_list)
+
+        return L1_align_feat_list
+    
+