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

spikezoo/archs/spikeformer/train.py

@@ -0,0 +1,173 @@
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = "0"
+import torch
+from torch import optim
+import numpy as np
+from DataProcess import DataLoader as dl
+from Model.SpikeFormer import SpikeFormer
+from Metrics.Metrics import Metrics
+from Model import Loss
+from utils import SaveModel, LoadModel
+from PIL import Image
+
+def eval(model, validData, epoch, optimizer, metrics):
+
+    model.eval()
+    print('Eval Epoch: %s' %(epoch))
+
+    with torch.no_grad():
+        pres = []
+        gts = []
+        for i, (spikes, gtImg) in enumerate(validData):
+
+            spikes = spikes.cuda()
+            gtImg = gtImg.cuda()
+            predImg = model(spikes)
+            predImg = predImg.squeeze(1)
+            predImg = predImg[:,3:-3,:]
+
+            predImg = predImg.clamp(min=-1., max=1.)
+            predImg = predImg.detach().cpu().numpy()
+            gtImg = gtImg.clamp(min=-1., max=1.)
+            gtImg = gtImg.detach().cpu().numpy()
+
+            predImg = (predImg + 1.) / 2. * 255.
+            predImg = predImg.astype(np.uint8)
+
+            gtImg = (gtImg + 1.) / 2. * 255.
+            gtImg = gtImg.astype(np.uint8)
+
+            pres.append(predImg)
+            gts.append(gtImg)
+        pres = np.concatenate(pres, axis=0)
+        gts = np.concatenate(gts, axis=0)
+
+        psnr = metrics.Cal_PSNR(pres, gts)
+        ssim = metrics.Cal_SSIM(pres, gts)
+        best_psnr, best_ssim, _ = metrics.GetBestMetrics()
+
+        SaveModel(epoch, (psnr, ssim), model, optimizer, saveRoot)
+        if psnr >= best_psnr and ssim >= best_ssim:
+            metrics.Update(psnr, ssim)
+            SaveModel(epoch, (psnr, ssim), model, optimizer, saveRoot, best=True)
+            with open('eval_best_log.txt', 'w') as f:
+                f.write('epoch: %s; psnr: %s, ssim: %s\n' %(epoch, psnr, ssim))
+            B, H, W = pres.shape
+            divide_line = np.zeros((H,4)).astype(np.uint8)
+            num = 0
+            for pre, gt in zip(pres, gts):
+                num += 1
+                concatImg = np.concatenate([pre, divide_line, gt], axis=1)
+                concatImg = Image.fromarray(concatImg)
+                concatImg.save('EvalResults/valid_%s.jpg' % (num))
+
+        print('*********************************************************')
+        best_psnr, best_ssim, _ = metrics.GetBestMetrics()
+        print('Eval Epoch: %s, PSNR: %s, SSIM: %s, Best_PSNR: %s, Best_SSIM: %s'
+              %(epoch, psnr, ssim, best_psnr, best_ssim))
+
+    model.train()
+
+def Train(trainData, validData, model, optimizer, epoch, start_epoch, metrics, saveRoot, perIter):
+    avg_l2_loss = 0.
+    avg_vgg_loss = 0.
+    avg_edge_loss = 0.
+    avg_total_loss = 0.
+    l2loss = Loss.CharbonnierLoss()
+    vggloss = Loss.VGGLoss4('vgg19-low-level4.pth').cuda()
+    criterion_edge = Loss.EdgeLoss()
+    LAMBDA_L2 = 100.0
+    LAMBDA_VGG = 1.0
+    LAMBDA_EDGE = 5.0   
+    for i in range(start_epoch, epoch):
+        for iter, (spikes, gtImg) in enumerate(trainData):
+            spikes = spikes.cuda()
+            gtImg = gtImg.cuda()
+            predImg = model(spikes)
+            gtImg = gtImg.unsqueeze(1)
+            predImg = predImg[:,:,3:-3,:]
+
+            loss_vgg = vggloss(gtImg, predImg) * LAMBDA_VGG
+            loss_l2 = l2loss(gtImg, predImg) * LAMBDA_L2
+            loss_edge = criterion_edge(gtImg, predImg) * LAMBDA_EDGE
+
+            totalLoss = loss_l2 + loss_vgg + loss_edge
+
+            optimizer.zero_grad()
+            totalLoss.backward()
+            optimizer.step()
+
+            avg_l2_loss += loss_l2.detach().cpu()
+            avg_vgg_loss += loss_vgg.detach().cpu()
+            avg_edge_loss += loss_edge.detach().cpu()
+            avg_total_loss += totalLoss.detach().cpu()
+            if (iter + 1) % perIter == 0:
+                avg_l2_loss = avg_l2_loss / perIter
+                avg_vgg_loss = avg_vgg_loss / perIter
+                avg_edge_loss = avg_edge_loss / perIter
+                avg_total_loss = avg_total_loss / perIter
+                print('=============================================================')
+                print('Epoch: %s, Iter: %s' % (i, iter + 1))
+                print('L2Loss: %s; VggLoss: %s; EdgeLoss: %s; TotalLoss: %s' % (
+                    avg_l2_loss.item(), avg_vgg_loss.item(), avg_edge_loss.item(), avg_total_loss.item()))
+                avg_l2_loss = 0.
+                avg_vgg_loss = 0.
+                avg_edge_loss = 0.
+                avg_total_loss = 0.
+
+        if (i + 1) % 1 == 0:
+            eval(model, validData, i, optimizer, metrics)
+
+if __name__ == "__main__":
+
+    dataPath = "/home/storage2/shechen/Spike_Sample_250x400"
+    spikeRadius = 32  # half length of input spike sequence expcept for the middle frame
+    spikeLen = 2 * spikeRadius + 1  # length of input spike sequence
+    batchSize = 2
+    epoch = 200
+    start_epoch = 0
+    lr = 2e-4
+    saveRoot = "CheckPoints/"  # path to save the trained model
+    perIter = 20
+
+    reuse = False
+    reuseType = 'latest'  # 'latest' or 'best'
+    checkPath = os.path.join('CheckPoints', '%s.pth' % (reuseType))
+
+    trainContainer = dl.DataContainer(dataPath=dataPath, dataType='train',
+                                      spikeRadius=spikeRadius,
+                                      batchSize=batchSize)
+    trainData = trainContainer.GetLoader()
+
+    validContainer = dl.DataContainer(dataPath=dataPath, dataType='valid',
+                                      spikeRadius=spikeRadius,
+                                      batchSize=batchSize)
+    validData = validContainer.GetLoader()
+
+    metrics = Metrics()
+
+    model = SpikeFormer(
+        inputDim = spikeLen,
+        dims = (32, 64, 160, 256),      # dimensions of each stage
+        heads = (1, 2, 5, 8),           # heads of each stage
+        ff_expansion = (8, 8, 4, 4),    # feedforward expansion factor of each stage
+        reduction_ratio = (8, 4, 2, 1), # reduction ratio of each stage for efficient attention
+        num_layers = 2,                 # num layers of each stage
+        decoder_dim = 256,              # decoder dimension
+        out_channel = 1                 # channel of restored image
+    ).cuda()
+
+    optimizer = optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.999), amsgrad=False)
+
+    if reuse:
+        preEpoch, prePerformance, modelDict, optDict = LoadModel(checkPath, model, optimizer)
+        start_epoch = preEpoch + 1
+        psnr, ssim = prePerformance[0], prePerformance[1]
+        metrics.Update(psnr, ssim)
+        for para in optimizer.param_groups:
+            para['lr'] = lr
+
+    model.train()
+
+    Train(trainData, validData, model, optimizer, epoch, start_epoch,
+                    metrics, saveRoot, perIter)

spikezoo/archs/spikeformer/utils.py

@@ -0,0 +1,22 @@
+import os
+import torch
+
+def SaveModel(epoch, bestPerformance, model, optimizer, saveRoot, best=False):
+    saveDict = {
+        'pre_epoch':epoch,
+        'performance':bestPerformance,
+        'model_state_dict':model.state_dict(),
+        'optimizer_state_dict':optimizer.state_dict()
+    }
+    savePath = os.path.join(saveRoot, '%s.pth' %('latest' if not best else 'best'))
+    torch.save(saveDict, savePath)
+
+def LoadModel(checkPath, model, optimizer=None):
+    stateDict = torch.load(checkPath)
+    pre_epoch = stateDict['pre_epoch']
+    model.load_state_dict(stateDict['model_state_dict'])
+    if optimizer is not None:
+        optimizer.load_state_dict(stateDict['optimizer_state_dict'])
+
+    return pre_epoch, stateDict['performance'], \
+           stateDict['model_state_dict'], stateDict['optimizer_state_dict']

spikezoo/archs/spk2imgnet/.github/workflows/pylint.yml

@@ -0,0 +1,23 @@
+name: Pylint
+
+on: [push]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.8", "3.9", "3.10"]
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v3
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install pylint
+    - name: Analysing the code with pylint
+      run: |
+        pylint $(git ls-files '*.py')

spikezoo/archs/spk2imgnet/.gitignore

@@ -0,0 +1,150 @@
+### Python template
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+.idea/
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+results/
+ckpt/
+ckpt2/
+old_ckpt/
+Spk2ImgNet_test2/
+Spk2ImgNet_train/
+*.zip
+*.pth
+*.h5
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+

spikezoo/archs/spk2imgnet/DCNv2.py

@@ -0,0 +1,135 @@
+#!/usr/bin/env python
+
+import math
+import logging
+import torch
+from PIL.Image import logger
+from torch import nn
+import torchvision
+from torch.nn.modules.utils import _pair
+
+
+class DCNv2(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation=1,
+                 deformable_groups=1):
+        super(DCNv2, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = _pair(kernel_size)
+        self.stride = _pair(stride)
+        self.padding = _pair(padding)
+        self.dilation = _pair(dilation)
+        self.deformable_groups = deformable_groups
+
+        self.weight = nn.Parameter(torch.Tensor(out_channels, in_channels, *self.kernel_size))
+        self.bias = nn.Parameter(torch.Tensor(out_channels))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        n = self.in_channels
+        for k in self.kernel_size:
+            n *= k
+        stdv = 1. / math.sqrt(n)
+        self.weight.data.uniform_(-stdv, stdv)
+        self.bias.data.zero_()
+
+    def forward(self, input, offset, mask):
+        assert 2 * self.deformable_groups * self.kernel_size[0] * self.kernel_size[1] == \
+            offset.shape[1]
+        assert self.deformable_groups * self.kernel_size[0] * self.kernel_size[1] == \
+            mask.shape[1]
+
+
+        return torchvision.ops.deform_conv2d(
+            input=input,
+            offset=offset,
+            mask=mask,
+            weight=self.weight,
+            bias=self.bias,
+            stride=self.stride,
+            padding=self.padding,
+            dilation=self.dilation,
+            groups=self.deformable_groups
+        )
+
+
+
+
+class DCN(DCNv2):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation=1,
+                 deformable_groups=1):
+        super(DCN, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation,
+                                  deformable_groups)
+
+        channels_ = self.deformable_groups * 3 * self.kernel_size[0] * self.kernel_size[1]
+        self.conv_offset_mask = nn.Conv2d(self.in_channels, channels_, kernel_size=self.kernel_size,
+                                          stride=self.stride, padding=self.padding, bias=True)
+        self.init_offset()
+
+    def init_offset(self):
+        self.conv_offset_mask.weight.data.zero_()
+        self.conv_offset_mask.bias.data.zero_()
+
+    def forward(self, input):
+        out = self.conv_offset_mask(input)
+        o1, o2, mask = torch.chunk(out, 3, dim=1)
+        offset = torch.cat((o1, o2), dim=1)
+        mask = torch.sigmoid(mask)
+
+        # return dcn_v2_conv(input, offset, mask, self.weight, self.bias, self.stride, self.padding,
+        #                    self.dilation, self.deformable_groups)
+
+        return torchvision.ops.deform_conv2d(
+            input=input,
+            offset=offset,
+            mask=mask,
+            weight=self.weight,
+            bias=self.bias,
+            stride=self.stride,
+            padding=self.padding,
+            dilation=self.dilation,
+        )
+
+
+class DCN_sep(DCNv2):
+    '''Use other features to generate offsets and masks'''
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation=1,
+                 deformable_groups=1):
+        super(DCN_sep, self).__init__(in_channels, out_channels, kernel_size, stride, padding,
+                                      dilation, deformable_groups)
+
+        channels_ = self.deformable_groups * 3 * self.kernel_size[0] * self.kernel_size[1]
+        self.conv_offset_mask = nn.Conv2d(self.in_channels, channels_, kernel_size=self.kernel_size,
+                                          stride=self.stride, padding=self.padding, bias=True)
+        self.init_offset()
+
+    def init_offset(self):
+        self.conv_offset_mask.weight.data.zero_()
+        self.conv_offset_mask.bias.data.zero_()
+
+    def forward(self, input, fea):
+        '''input: input features for deformable conv
+        fea: other features used for generating offsets and mask'''
+        out = self.conv_offset_mask(fea)
+        o1, o2, mask = torch.chunk(out, 3, dim=1)
+        offset = torch.cat((o1, o2), dim=1)
+
+        offset_mean = torch.mean(torch.abs(offset))
+        if offset_mean > 100:
+            logger.warning('Offset mean is {}, larger than 100.'.format(offset_mean))
+
+        mask = torch.sigmoid(mask)
+
+
+        return torchvision.ops.deform_conv2d(
+            input=input,
+            offset=offset,
+            mask=mask,
+            weight=self.weight,
+            bias=self.bias,
+            stride=self.stride,
+            padding=self.padding,
+            dilation=self.dilation,
+        )
+

spikezoo/archs/spk2imgnet/align_arch.py

@@ -0,0 +1,159 @@
+""" network architecture for Sakuya """
+import torch
+from DCNv2 import *
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.ops import DeformConv2d
+
+
+
+
+class PCDAlign(nn.Module):
+    """Alignment module using Pyramid, Cascading and Deformable convolution
+    with 3 pyramid levels.
+    """
+
+    def __init__(self, nf=64, groups=8):
+        super(PCDAlign, self).__init__()
+
+        # fea1
+        # L3: level 3, 1/4 spatial size
+        self.L3_offset_conv1_1 = nn.Conv2d(
+            nf * 2, nf, 3, 1, 1, bias=True
+        )  # concat for dif
+        self.L3_offset_conv2_1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.L3_dcnpack_1 = DCN_sep(
+            nf, nf, 3, stride=1, padding=1, dilation=1, deformable_groups=groups
+        )
+        # L2: level 2, 1/2 spatial size
+        self.L2_offset_conv1_1 = nn.Conv2d(
+            nf * 2, nf, 3, 1, 1, bias=True
+        )  # concat for diff
+        self.L2_offset_conv2_1 = nn.Conv2d(
+            nf * 2, nf, 3, 1, 1, bias=True
+        )  # concat for offset
+        self.L2_offset_conv3_1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.L2_dcnpack_1 = DCN_sep(
+            nf, nf, 3, stride=1, padding=1, dilation=1, deformable_groups=groups
+        )
+        self.L2_fea_conv_1 = nn.Conv2d(nf * 2, nf, 3, 1, 1, bias=True)  # concat for fea
+        # L1: level 1, original spatial size
+        self.L1_offset_conv1_1 = nn.Conv2d(
+            nf * 2, nf, 3, 1, 1, bias=True
+        )  # concat for diff
+        self.L1_offset_conv2_1 = nn.Conv2d(
+            nf * 2, nf, 3, 1, 1, bias=True
+        )  # concat for offset
+        self.L1_offset_conv3_1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.L1_dcnpack_1 = DCN_sep(
+            nf, nf, 3, stride=1, padding=1, dilation=1, deformable_groups=groups
+        )
+        self.L1_fea_conv_1 = nn.Conv2d(nf * 2, nf, 3, 1, 1, bias=True)  # concat for fea
+
+        # Cascading DCN
+        self.cas_offset_conv1 = nn.Conv2d(
+            nf * 2, nf, 3, 1, 1, bias=True
+        )  # concat for diff
+        self.cas_offset_conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.cas_dcnpack = DCN_sep(
+            nf, nf, 3, stride=1, padding=1, dilation=1, deformable_groups=groups
+        )
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+
+    def forward(self, fea1, fea2):
+        """align other neighboring frames to the reference frame in the feature level
+        fea1, fea2: [L1, L2, L3], each with [B,C,H,W] features
+        fea1 : features of neighboring frame
+        fea2 : features of reference (key) frame
+        estimate offset bidirectionally
+        """
+        # param. of fea1
+        # L3
+        L3_offset = torch.cat([fea1[2], fea2[2]], dim=1)
+        L3_offset = self.lrelu(self.L3_offset_conv1_1(L3_offset))
+        L3_offset = self.lrelu(self.L3_offset_conv2_1(L3_offset))
+        L3_fea = self.lrelu(self.L3_dcnpack_1(fea1[2], L3_offset))
+        # L2
+        L2_offset = torch.cat([fea1[1], fea2[1]], dim=1)
+        L2_offset = self.lrelu(self.L2_offset_conv1_1(L2_offset))
+        L3_offset = F.interpolate(
+            L3_offset, scale_factor=2, mode="bilinear", align_corners=False
+        )
+        L2_offset = self.lrelu(
+            self.L2_offset_conv2_1(torch.cat([L2_offset, L3_offset * 2], dim=1))
+        )
+        L2_offset = self.lrelu(self.L2_offset_conv3_1(L2_offset))
+        L2_fea = self.L2_dcnpack_1(fea1[1], L2_offset)
+        L3_fea = F.interpolate(
+            L3_fea, scale_factor=2, mode="bilinear", align_corners=False
+        )
+        L2_fea = self.lrelu(self.L2_fea_conv_1(torch.cat([L2_fea, L3_fea], dim=1)))
+        # L1
+        L1_offset = torch.cat([fea1[0], fea2[0]], dim=1)
+        L1_offset = self.lrelu(self.L1_offset_conv1_1(L1_offset))
+        L2_offset = F.interpolate(
+            L2_offset, scale_factor=2, mode="bilinear", align_corners=False
+        )
+        L1_offset = self.lrelu(
+            self.L1_offset_conv2_1(torch.cat([L1_offset, L2_offset * 2], dim=1))
+        )
+        L1_offset = self.lrelu(self.L1_offset_conv3_1(L1_offset))
+        L1_fea = self.L1_dcnpack_1(fea1[0], L1_offset)
+        L2_fea = F.interpolate(
+            L2_fea, scale_factor=2, mode="bilinear", align_corners=False
+        )
+        L1_fea = self.L1_fea_conv_1(torch.cat([L1_fea, L2_fea], dim=1))
+
+        # Cascading DCN
+        offset = torch.cat([L1_fea, fea2[0]], dim=1)
+        offset = self.lrelu(self.cas_offset_conv1(offset))
+        offset = self.lrelu(self.cas_offset_conv2(offset))
+        L1_fea = self.lrelu(self.cas_dcnpack(L1_fea, offset))
+
+        return L1_fea
+
+
+class Easy_PCD(nn.Module):
+    def __init__(self, nf=64, groups=8):
+        super(Easy_PCD, self).__init__()
+
+        self.fea_L2_conv1 = nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
+        self.fea_L2_conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.fea_L3_conv1 = nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
+        self.fea_L3_conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+        self.pcd_align = PCDAlign(nf=nf, groups=groups)
+        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+
+    def forward(self, f1, f2):
+        # input: extracted features
+        # f1: feature of neighboring frame
+        # f2: feature of the key (reference) frame
+        # feature size: f1 = f2 = [B, C, H, W]
+        # print(f1.size())
+        L1_fea = torch.stack([f1, f2], dim=1)  # [B, 2, C, H, W]
+        B, N, C, H, W = L1_fea.size()
+        L1_fea = L1_fea.view(-1, C, H, W)
+        # L2
+        L2_fea = self.lrelu(self.fea_L2_conv1(L1_fea))
+        L2_fea = self.lrelu(self.fea_L2_conv2(L2_fea))
+        # L3
+        L3_fea = self.lrelu(self.fea_L3_conv1(L2_fea))
+        L3_fea = self.lrelu(self.fea_L3_conv2(L3_fea))
+
+        L1_fea = L1_fea.view(B, N, -1, H, W)
+        L2_fea = L2_fea.view(B, N, -1, H // 2, W // 2)
+        L3_fea = L3_fea.view(B, N, -1, H // 4, W // 4)
+
+        fea1 = [
+            L1_fea[:, 0, :, :, :].clone(),
+            L2_fea[:, 0, :, :, :].clone(),
+            L3_fea[:, 0, :, :, :].clone(),
+        ]
+        fea2 = [
+            L1_fea[:, 1, :, :, :].clone(),
+            L2_fea[:, 1, :, :, :].clone(),
+            L3_fea[:, 1, :, :, :].clone(),
+        ]
+        aligned_fea = self.pcd_align(fea1, fea2)
+        return aligned_fea