视频超分经典文章（一）EDVR: Video Restoration with Enhanced Deformable Convolutional Networks (包含代码)

摘要：

视频恢复任务，包括超分辨率、去模糊等，越来越受到计算机视觉界的关注。在NTIRE19挑战赛中发布了一个名为REDS的具有挑战性的基准测试。这个新的基准从两个方面挑战了现有的方法:(1)如何对给定大运动的多帧图像进行对齐，(2)如何有效地融合不同运动和模糊的不同帧。在这项工作中，我们提出了一个新的具有增强可变形卷积的视频恢复框架，称为EDVR，以解决这些挑战。首先，为了处理大的运动，我们设计了一个金字塔，级联和可变形(PCD)对齐模块，其中的帧对齐是在特征级使用可变形的卷积以粗到细的方式完成的。其次，我们提出了时间和空间注意力(TSA)融合模块，该模块将注意力应用于时间和空间上，以强调后续恢复的重要特征。由于这些模块，我们的EDVR赢得冠军，并在NTIRE19视频恢复和增强挑战的所有四个轨道上以巨大的优势超过第二名。EDVR在视频超分辨率和去模糊方面也显示出了优于已发表的最先进的方法的性能。

1.介绍

对齐。大多数现有方法通过明确估计参考帧与其相邻帧之间的光流场来执行对齐[2、48、13]。相邻帧根据估计的运动场进行扭曲。另一类研究通过动态滤波[10]或可变形卷积[40]实现隐式运动补偿。REDS对现有的对齐算法提出了巨大的挑战。特别是，对于基于流的方法来说，精确的流估计和准确的扭曲可能具有挑战性且耗时。在存在大运动的情况下，难以在单一分辨率尺度内明确或隐式地进行运动补偿。

融合。将对齐帧的特征进行融合是视频修复任务中的另一个关键步骤。大多数现有方法要么使用卷积在所有帧上进行早期融合[2]，要么采用递归网络逐渐融合多个帧[32、6]。刘等人[22]提出了一个时间自适应网络，可以在不同的时间尺度上动态融合。这些现有方法都没有考虑每个帧的底层视觉信息的重要性——不同的帧和位置对重建的贡献并不相等，因为一些帧或区域受到了不完美的对齐和模糊的影响。

解决方案。提出了一个统一的框架，称为EDVR，可扩展到各种视频修复任务，包括超分辨率和去模糊。EDVR的核心包括：（1）一个称为金字塔级联和可变形卷积（PCD）的对齐模块，以及（2）一个称为时空注意力（TSA）的融合模块。

PCD模块受到TDAN[40]的启发，使用可变形卷积在特征级别上将每个相邻帧与参考帧对齐。与TDAN不同的是，采用粗到精的方式进行对齐，以处理大规模和复杂的运动。具体而言，我们使用金字塔结构，首先使用粗略的估计将较低尺度的特征对齐，然后将偏移量和对齐特征传播到较高尺度，以促进精确的运动补偿，类似于光流估计中采用的概念[7, 9]。此外，在金字塔对齐操作之后，我们级联了一个额外的可变形卷积，进一步提高了对齐的鲁棒性。

提出的TSA是一个融合模块，有助于跨多个对齐特征聚合信息。为了更好地考虑每个帧的视觉信息，通过计算参考帧和每个相邻帧的特征之间的逐元素相关性引入了时间注意力。然后，在每个位置上，相关系数对每个相邻特征进行加权，表示其对重建参考图像的信息量。然后将所有帧的加权特征进行卷积和融合。在具有时间注意力的融合之后，我们进一步应用空间注意力，对每个通道中的每个位置分配权重，以更有效地利用跨通道和空间信息。

我们参加了视频修复和增强挑战的四个赛道[29, 28]，包括视频超分辨率（清晰/模糊）和视频去模糊（清晰/压缩伪影）。由于有效的对齐和融合模块，我们的EDVR在所有四个具有挑战性的赛道中获得了冠军，展示了我们方法的效果和通用性。除了比赛结果，我们还在现有的视频超分辨率和去模糊基准测试中报告了比较结果。在这些视频修复任务中，我们的EDVR表现出优于现有方法的性能。

1.1 可变形卷积

可变形卷积（Deformable Convolution）是一种卷积神经网络中的操作，它允许网络在进行卷积时对输入特征图进行局部的空间变形。传统的卷积操作只能在固定窗口内进行卷积计算，而可变形卷积则引入了额外的偏移量（offsets）来动态调整局部感受野的采样位置。

在可变形卷积中，每个位置的卷积核采样位置是通过应用一个偏移量来计算的。这些偏移量是在训练过程中通过学习得到的，它们表示了输入特征图中每个位置的局部偏移信息。通过使用这些偏移量，可变形卷积可以根据输入特征图的局部结构自适应地调整卷积核的采样位置，从而更好地适应图像中的变形和不规则模式。

可以看到可变形卷积在普通卷积的基础上加了一个偏移量，偏移量是由另一个卷积生成的，通常是小数。因为是小数，所以对不上实际的像素点，这里需要通过双线性插值来确定该点的像素值。 

优点：目标物体往往具有不同尺度的大小，但是传统的CNN只具有固定的感受野，这就导致特征提取效果不佳。而可变形卷积允许感受野内的像素学习一个偏移量自适应地去调整采样点的位置，这能提高模型的特征提取能力和鲁棒性。

2.模型

2.1总体结构

以视频超分辨率为例，EDVR将2N+1个低分辨率帧作为输入，并生成一个高分辨率的输出。每个相邻帧都通过PCD对齐模块在特征级别上与参考帧对齐。TSA融合模块融合不同帧的图像信息。这两个模块的详细信息在第3.2节和第3.3节中进行了描述。融合后的特征经过重建模块，该模块是EDVR中一系列残差块的级联，可以轻松地用单幅图像超分辨率中的任何其他先进模块进行替换[46, 51]。上采样操作在网络的末端进行，以增加空间尺寸。最后，通过将预测的图像残差添加到直接上采样的图像上，得到高分辨率帧。

对于其他具有高空间分辨率输入的任务，比如视频去模糊，输入帧首先通过步幅卷积层进行下采样。然后，大部分计算都在低分辨率空间中进行，这在很大程度上节省了计算成本。末端的上采样层将特征调整回原始输入分辨率。在对齐模块之前使用了一个预去模糊模块，对模糊输入进行预处理，提高对齐精度。

2.2 使用金字塔、级联和可变形卷积进行对齐

我们首先简要回顾了用于对齐[40]的可变形卷积的使用，即，对齐每个相邻帧的特征到参考帧的特征。与基于光流的方法不同，该方法对每一帧的特征进行变形对齐，用Ft+i, i∈[−N:+N]表示。我们使用调制的可变形模块[53]。给定K个采样位置的可变形卷积核，我们将wk和pk分别表示为第K个位置的权值和预定的偏移量。例如，3×3内核定义为K=9, pk∈{(−1，−1)，(−1,0)，···，(1,1)}。在每个位置p0处的对齐特征Fa t+i可以通过以下方法得到:

其中，∆P={∆P}， f为由多个卷积层组成的一般函数，[.,.]为串联运算。为简单起见，在描述和图形中，我们只考虑∆pk的可学习偏移，忽略∆mk的调制。由于p0 +pk +∆pk为分数，采用[3]中的双线性插值。

为了解决对齐过程中的复杂运动和大视差问题，我们基于光流中成熟的原理提出了PCD模块:金字塔处理[31,35]和级联细化[7,8,9]。具体来说，如图3中黑色虚线所示，为了生成第l层的特征Fl t+i，我们使用跨步卷积滤波器将第(l−1)层的特征降采样2倍，得到特征表示的l层金字塔。在第l层，偏移量和对齐特征也分别由第l+1层的上采样偏移量和对齐特征预测(图3中的紫色虚线):

式中(·)↑s表示按因子s进行缩放，DConv为eq . 1中描述的可变形卷积，g为具有多个卷积层的一般函数。×2上采样采用双线性插值实现。我们在EDVR中使用三层金字塔，即L=3。为了降低计算成本，我们不随着空间大小的减小而增加信道数。

在金字塔结构的基础上，进行后续的可变形对齐，进一步细化粗对齐的特征(图3中背景为浅紫色的部分)。这种由粗到细的方式使得PCD模块将对齐提高到亚像素精度。我们在4.3节中论证了PCD的有效性。值得注意的是，PCD对齐模块是与整个框架一起共同学习的，没有额外的监督[40]，也没有光流[48]等其他任务的预训练。 

2.3时间和空间注意力的融合

帧间时间关系和帧内空间关系是融合的关键因为1)由于遮挡、模糊区域和视差问题，不同的相邻帧信息不相等;2)前一对准阶段产生的不对准和不对准对后续重建性能产生不利影响。因此，对相邻帧进行像素级动态聚合是实现有效融合的必要条件。为了解决上述问题，我们提出了TSA融合模块来为每一帧分配像素级聚合权值。具体来说，我们在融合过程中采用了时间和空间的注意，如图4时间注意力的目标是在嵌入空间中计算帧的相似度。直观地说，在一个嵌入空间中，应该更多地注意一个更类似于参考帧的相邻帧。对于每一帧i∈[−N:+N]，相似距离h可计算为:

 其中θ(Fa t+i)和φ(Fa t)是两种嵌入，可以通过简单的卷积滤波器实现。使用sigmoid激活函数限制输出在[0,1]，稳定梯度反向传播。注意，对于每个空间位置，时间注意都具有空间特异性，即h(Fa t+i, Fa t)的空间大小与Fa t+i相同。

然后将时间注意图以像素方式乘以原始对齐特征Fa t+i。额外的融合卷积层被采用来聚合这些注意力调制的特征Fa t+i:

其中⊙和[·，·，·]分别表示元素级乘法和级联。

然后通过融合特征计算出空间注意掩模。采用金字塔设计，增加注意接受场。然后，通过掩模的逐元乘法和加法对融合特征进行调制，类似于[45]。TSA模块的有效性见第4.3节。

2.4 二阶段重建 

虽然单个EDVR装备PCD对齐模块和TSA融合模块可以达到目前最先进的性能，但观察到恢复后的图像偏置偏置级联DConv金字塔并不完美，特别是当输入帧模糊或严重失真时。在这种恶劣的环境下，运动补偿和细节聚合受到影响，导致重建性能较差。

直观地说，粗还原框架将大大减轻对齐和融合的压力。因此，我们采用两阶段策略来进一步提高性能。具体来说，一个类似但较浅的EDVR网络被级联，以细化第一阶段的输出帧。好处有两方面:1)有效消除了前一模型无法处理的严重运动模糊，提高了恢复质量;2)缓解了输出帧之间的不一致性。

3.实验

3.1 实验数据

训练数据集。以往关于视频处理的研究[21,10,34]通常是在私有数据集上发展或评估的。缺乏标准和开放的视频数据集限制了公平的比较。REDS[26]是NTIRE19竞赛中新提出的高质量(720p)视频数据集。REDS由240个训练片段、30个验证片段和30个测试片段组成(每个片段有100个连续帧)。在比赛期间，由于没有测试场地的真实情况，我们选择了4个具有代表性的片段(场景和动作多样)作为我们的测试集，标记为REDS41。剩下的训练和验证片段被重新分组为我们的训练数据集(共有266个片段)。为了与我们在竞争中的方法和过程相一致，本文也采用了这种配置。

Vimeo-90K[48]是一个广泛使用的训练数据集，通常会与Vid4[21]和Vimeo-90K测试数据集(记为Vimeo-90K- t)一起进行评估。当训练集的分布与测试集的分布偏离时，我们观察到数据集的偏差。

训练细节。PCD对齐模块采用5个残差块(RB)进行特征提取。我们在重建模块中使用40个RBs，在第二阶段模型中使用20个RBs。每个残留块的通道大小设置为128。我们分别使用大小为64×64和256×256的RGB补丁作为视频SR和去模糊任务的输入。迷你批量大小设置为32。网络采用5个连续帧(即N=2)作为输入，除非另有规定。我们通过随机水平翻转和90◦旋转来增加训练数据。我们只采用Charbonnier惩罚函数[17]作为最终损失，定义为：

通过设置β1=0.9和β2=0.999，我们使用Adam优化器[14]训练我们的模型。初始化学习速率为4×10−4。我们用浅层网络的参数来初始化更深层次的网络，以便更快地收敛。我们使用PyTorch框架来实现我们的模型，并使用8个NVIDIA Titan Xp gpu来训练它们。 

4.代码

import torch
from torch import nn as nn
from torch.nn import functional as F

from basicsr.models.archs.arch_util import (DCNv2Pack, ResidualBlockNoBN,
                                            make_layer)


class PCDAlignment(nn.Module):
    """Alignment module using Pyramid, Cascading and Deformable convolution
    (PCD). It is used in EDVR.

    Ref:
        EDVR: Video Restoration with Enhanced Deformable Convolutional Networks

    Args:
        num_feat (int): Channel number of middle features. Default: 64.
        deformable_groups (int): Deformable groups. Defaults: 8.
    """

    def __init__(self, num_feat=64, deformable_groups=8):
        super(PCDAlignment, self).__init__()

        # Pyramid has three levels:
        # L3: level 3, 1/4 spatial size
        # L2: level 2, 1/2 spatial size
        # L1: level 1, original spatial size
        self.offset_conv1 = nn.ModuleDict()
        self.offset_conv2 = nn.ModuleDict()
        self.offset_conv3 = nn.ModuleDict()
        self.dcn_pack = nn.ModuleDict()
        self.feat_conv = nn.ModuleDict()

        # Pyramids
        for i in range(3, 0, -1):
            level = f'l{i}'
            self.offset_conv1[level] = nn.Conv2d(num_feat * 2, num_feat, 3, 1,
                                                 1)
            if i == 3:
                self.offset_conv2[level] = nn.Conv2d(num_feat, num_feat, 3, 1,
                                                     1)
            else:
                self.offset_conv2[level] = nn.Conv2d(num_feat * 2, num_feat, 3,
                                                     1, 1)
                self.offset_conv3[level] = nn.Conv2d(num_feat, num_feat, 3, 1,
                                                     1)
            self.dcn_pack[level] = DCNv2Pack(
                num_feat,
                num_feat,
                3,
                padding=1,
                deformable_groups=deformable_groups)

            if i < 3:
                self.feat_conv[level] = nn.Conv2d(num_feat * 2, num_feat, 3, 1,
                                                  1)

        # Cascading dcn
        self.cas_offset_conv1 = nn.Conv2d(num_feat * 2, num_feat, 3, 1, 1)
        self.cas_offset_conv2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.cas_dcnpack = DCNv2Pack(
            num_feat,
            num_feat,
            3,
            padding=1,
            deformable_groups=deformable_groups)

        self.upsample = nn.Upsample(
            scale_factor=2, mode='bilinear', align_corners=False)
        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)

    def forward(self, nbr_feat_l, ref_feat_l):
        """Align neighboring frame features to the reference frame features.

        Args:
            nbr_feat_l (list[Tensor]): Neighboring feature list. It
                contains three pyramid levels (L1, L2, L3),
                each with shape (b, c, h, w).
            ref_feat_l (list[Tensor]): Reference feature list. It
                contains three pyramid levels (L1, L2, L3),
                each with shape (b, c, h, w).

        Returns:
            Tensor: Aligned features.
        """
        # Pyramids
        upsampled_offset, upsampled_feat = None, None
        for i in range(3, 0, -1):
            level = f'l{i}'
            offset = torch.cat([nbr_feat_l[i - 1], ref_feat_l[i - 1]], dim=1)
            offset = self.lrelu(self.offset_conv1[level](offset))
            if i == 3:
                offset = self.lrelu(self.offset_conv2[level](offset))
            else:
                offset = self.lrelu(self.offset_conv2[level](torch.cat(
                    [offset, upsampled_offset], dim=1)))
                offset = self.lrelu(self.offset_conv3[level](offset))

            feat = self.dcn_pack[level](nbr_feat_l[i - 1], offset)
            if i < 3:
                feat = self.feat_conv[level](
                    torch.cat([feat, upsampled_feat], dim=1))
            if i > 1:
                feat = self.lrelu(feat)

            if i > 1:  # upsample offset and features
                # x2: when we upsample the offset, we should also enlarge
                # the magnitude.
                upsampled_offset = self.upsample(offset) * 2
                upsampled_feat = self.upsample(feat)

        # Cascading
        offset = torch.cat([feat, ref_feat_l[0]], dim=1)
        offset = self.lrelu(
            self.cas_offset_conv2(self.lrelu(self.cas_offset_conv1(offset))))
        feat = self.lrelu(self.cas_dcnpack(feat, offset))
        return feat


class TSAFusion(nn.Module):
    """Temporal Spatial Attention (TSA) fusion module.

    Temporal: Calculate the correlation between center frame and
        neighboring frames;
    Spatial: It has 3 pyramid levels, the attention is similar to SFT.
        (SFT: Recovering realistic texture in image super-resolution by deep
            spatial feature transform.)

    Args:
        num_feat (int): Channel number of middle features. Default: 64.
        num_frame (int): Number of frames. Default: 5.
        center_frame_idx (int): The index of center frame. Default: 2.
    """

    def __init__(self, num_feat=64, num_frame=5, center_frame_idx=2):
        super(TSAFusion, self).__init__()
        self.center_frame_idx = center_frame_idx
        # temporal attention (before fusion conv)
        self.temporal_attn1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.temporal_attn2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.feat_fusion = nn.Conv2d(num_frame * num_feat, num_feat, 1, 1)

        # spatial attention (after fusion conv)
        self.max_pool = nn.MaxPool2d(3, stride=2, padding=1)
        self.avg_pool = nn.AvgPool2d(3, stride=2, padding=1)
        self.spatial_attn1 = nn.Conv2d(num_frame * num_feat, num_feat, 1)
        self.spatial_attn2 = nn.Conv2d(num_feat * 2, num_feat, 1)
        self.spatial_attn3 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.spatial_attn4 = nn.Conv2d(num_feat, num_feat, 1)
        self.spatial_attn5 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.spatial_attn_l1 = nn.Conv2d(num_feat, num_feat, 1)
        self.spatial_attn_l2 = nn.Conv2d(num_feat * 2, num_feat, 3, 1, 1)
        self.spatial_attn_l3 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.spatial_attn_add1 = nn.Conv2d(num_feat, num_feat, 1)
        self.spatial_attn_add2 = nn.Conv2d(num_feat, num_feat, 1)

        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
        self.upsample = nn.Upsample(
            scale_factor=2, mode='bilinear', align_corners=False)

    def forward(self, aligned_feat):
        """
        Args:
            aligned_feat (Tensor): Aligned features with shape (b, t, c, h, w).

        Returns:
            Tensor: Features after TSA with the shape (b, c, h, w).
        """
        b, t, c, h, w = aligned_feat.size()
        # temporal attention
        embedding_ref = self.temporal_attn1(
            aligned_feat[:, self.center_frame_idx, :, :, :].clone())
        embedding = self.temporal_attn2(aligned_feat.view(-1, c, h, w))
        embedding = embedding.view(b, t, -1, h, w)  # (b, t, c, h, w)

        corr_l = []  # correlation list
        for i in range(t):
            emb_neighbor = embedding[:, i, :, :, :]
            corr = torch.sum(emb_neighbor * embedding_ref, 1)  # (b, h, w)
            corr_l.append(corr.unsqueeze(1))  # (b, 1, h, w)
        corr_prob = torch.sigmoid(torch.cat(corr_l, dim=1))  # (b, t, h, w)
        corr_prob = corr_prob.unsqueeze(2).expand(b, t, c, h, w)
        corr_prob = corr_prob.contiguous().view(b, -1, h, w)  # (b, t*c, h, w)
        aligned_feat = aligned_feat.view(b, -1, h, w) * corr_prob

        # fusion
        feat = self.lrelu(self.feat_fusion(aligned_feat))

        # spatial attention
        attn = self.lrelu(self.spatial_attn1(aligned_feat))
        attn_max = self.max_pool(attn)
        attn_avg = self.avg_pool(attn)
        attn = self.lrelu(
            self.spatial_attn2(torch.cat([attn_max, attn_avg], dim=1)))
        # pyramid levels
        attn_level = self.lrelu(self.spatial_attn_l1(attn))
        attn_max = self.max_pool(attn_level)
        attn_avg = self.avg_pool(attn_level)
        attn_level = self.lrelu(
            self.spatial_attn_l2(torch.cat([attn_max, attn_avg], dim=1)))
        attn_level = self.lrelu(self.spatial_attn_l3(attn_level))
        attn_level = self.upsample(attn_level)

        attn = self.lrelu(self.spatial_attn3(attn)) + attn_level
        attn = self.lrelu(self.spatial_attn4(attn))
        attn = self.upsample(attn)
        attn = self.spatial_attn5(attn)
        attn_add = self.spatial_attn_add2(
            self.lrelu(self.spatial_attn_add1(attn)))
        attn = torch.sigmoid(attn)

        # after initialization, * 2 makes (attn * 2) to be close to 1.
        feat = feat * attn * 2 + attn_add
        return feat


class PredeblurModule(nn.Module):
    """Pre-dublur module.

    Args:
        num_in_ch (int): Channel number of input image. Default: 3.
        num_feat (int): Channel number of intermediate features. Default: 64.
        hr_in (bool): Whether the input has high resolution. Default: False.
    """

    def __init__(self, num_in_ch=3, num_feat=64, hr_in=False):
        super(PredeblurModule, self).__init__()
        self.hr_in = hr_in

        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
        if self.hr_in:
            # downsample x4 by stride conv
            self.stride_conv_hr1 = nn.Conv2d(num_feat, num_feat, 3, 2, 1)
            self.stride_conv_hr2 = nn.Conv2d(num_feat, num_feat, 3, 2, 1)

        # generate feature pyramid
        self.stride_conv_l2 = nn.Conv2d(num_feat, num_feat, 3, 2, 1)
        self.stride_conv_l3 = nn.Conv2d(num_feat, num_feat, 3, 2, 1)

        self.resblock_l3 = ResidualBlockNoBN(num_feat=num_feat)
        self.resblock_l2_1 = ResidualBlockNoBN(num_feat=num_feat)
        self.resblock_l2_2 = ResidualBlockNoBN(num_feat=num_feat)
        self.resblock_l1 = nn.ModuleList(
            [ResidualBlockNoBN(num_feat=num_feat) for i in range(5)])

        self.upsample = nn.Upsample(
            scale_factor=2, mode='bilinear', align_corners=False)
        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)

    def forward(self, x):
        feat_l1 = self.lrelu(self.conv_first(x))
        if self.hr_in:
            feat_l1 = self.lrelu(self.stride_conv_hr1(feat_l1))
            feat_l1 = self.lrelu(self.stride_conv_hr2(feat_l1))

        # generate feature pyramid
        feat_l2 = self.lrelu(self.stride_conv_l2(feat_l1))
        feat_l3 = self.lrelu(self.stride_conv_l3(feat_l2))

        feat_l3 = self.upsample(self.resblock_l3(feat_l3))
        feat_l2 = self.resblock_l2_1(feat_l2) + feat_l3
        feat_l2 = self.upsample(self.resblock_l2_2(feat_l2))

        for i in range(2):
            feat_l1 = self.resblock_l1[i](feat_l1)
        feat_l1 = feat_l1 + feat_l2
        for i in range(2, 5):
            feat_l1 = self.resblock_l1[i](feat_l1)
        return feat_l1


class EDVR(nn.Module):
    """EDVR network structure for video super-resolution.

    Now only support X4 upsampling factor.
    Paper:
        EDVR: Video Restoration with Enhanced Deformable Convolutional Networks

    Args:
        num_in_ch (int): Channel number of input image. Default: 3.
        num_out_ch (int): Channel number of output image. Default: 3.
        num_feat (int): Channel number of intermediate features. Default: 64.
        num_frame (int): Number of input frames. Default: 5.
        deformable_groups (int): Deformable groups. Defaults: 8.
        num_extract_block (int): Number of blocks for feature extraction.
            Default: 5.
        num_reconstruct_block (int): Number of blocks for reconstruction.
            Default: 10.
        center_frame_idx (int): The index of center frame. Frame counting from
            0. Default: 2.
        hr_in (bool): Whether the input has high resolution. Default: False.
        with_predeblur (bool): Whether has predeblur module.
            Default: False.
        with_tsa (bool): Whether has TSA module. Default: True.
    """

    def __init__(self,
                 num_in_ch=3,
                 num_out_ch=3,
                 num_feat=64,
                 num_frame=5,
                 deformable_groups=8,
                 num_extract_block=5,
                 num_reconstruct_block=10,
                 center_frame_idx=2,
                 hr_in=False,
                 with_predeblur=False,
                 with_tsa=True):
        super(EDVR, self).__init__()
        if center_frame_idx is None:
            self.center_frame_idx = num_frame // 2
        else:
            self.center_frame_idx = center_frame_idx
        self.hr_in = hr_in
        self.with_predeblur = with_predeblur
        self.with_tsa = with_tsa

        # extract features for each frame
        if self.with_predeblur:
            self.predeblur = PredeblurModule(
                num_feat=num_feat, hr_in=self.hr_in)
            self.conv_1x1 = nn.Conv2d(num_feat, num_feat, 1, 1)
        else:
            self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)

        # extrat pyramid features
        self.feature_extraction = make_layer(
            ResidualBlockNoBN, num_extract_block, num_feat=num_feat)
        self.conv_l2_1 = nn.Conv2d(num_feat, num_feat, 3, 2, 1)
        self.conv_l2_2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.conv_l3_1 = nn.Conv2d(num_feat, num_feat, 3, 2, 1)
        self.conv_l3_2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)

        # pcd and tsa module
        self.pcd_align = PCDAlignment(
            num_feat=num_feat, deformable_groups=deformable_groups)
        if self.with_tsa:
            self.fusion = TSAFusion(
                num_feat=num_feat,
                num_frame=num_frame,
                center_frame_idx=self.center_frame_idx)
        else:
            self.fusion = nn.Conv2d(num_frame * num_feat, num_feat, 1, 1)

        # reconstruction
        self.reconstruction = make_layer(
            ResidualBlockNoBN, num_reconstruct_block, num_feat=num_feat)
        # upsample
        self.upconv1 = nn.Conv2d(num_feat, num_feat * 4, 3, 1, 1)
        self.upconv2 = nn.Conv2d(num_feat, 64 * 4, 3, 1, 1)
        self.pixel_shuffle = nn.PixelShuffle(2)
        self.conv_hr = nn.Conv2d(64, 64, 3, 1, 1)
        self.conv_last = nn.Conv2d(64, 3, 3, 1, 1)

        # activation function
        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)

    def forward(self, x):
        b, t, c, h, w = x.size()
        if self.hr_in:
            assert h % 16 == 0 and w % 16 == 0, (
                'The height and width must be multiple of 16.')
        else:
            assert h % 4 == 0 and w % 4 == 0, (
                'The height and width must be multiple of 4.')

        x_center = x[:, self.center_frame_idx, :, :, :].contiguous()

        # extract features for each frame
        # L1
        if self.with_predeblur:
            feat_l1 = self.conv_1x1(self.predeblur(x.view(-1, c, h, w)))
            if self.hr_in:
                h, w = h // 4, w // 4
        else:
            feat_l1 = self.lrelu(self.conv_first(x.view(-1, c, h, w)))

        feat_l1 = self.feature_extraction(feat_l1)
        # L2
        feat_l2 = self.lrelu(self.conv_l2_1(feat_l1))
        feat_l2 = self.lrelu(self.conv_l2_2(feat_l2))
        # L3
        feat_l3 = self.lrelu(self.conv_l3_1(feat_l2))
        feat_l3 = self.lrelu(self.conv_l3_2(feat_l3))

        feat_l1 = feat_l1.view(b, t, -1, h, w)
        feat_l2 = feat_l2.view(b, t, -1, h // 2, w // 2)
        feat_l3 = feat_l3.view(b, t, -1, h // 4, w // 4)

        # PCD alignment
        ref_feat_l = [  # reference feature list
            feat_l1[:, self.center_frame_idx, :, :, :].clone(),
            feat_l2[:, self.center_frame_idx, :, :, :].clone(),
            feat_l3[:, self.center_frame_idx, :, :, :].clone()
        ]
        aligned_feat = []
        for i in range(t):
            nbr_feat_l = [  # neighboring feature list
                feat_l1[:, i, :, :, :].clone(), feat_l2[:, i, :, :, :].clone(),
                feat_l3[:, i, :, :, :].clone()
            ]
            aligned_feat.append(self.pcd_align(nbr_feat_l, ref_feat_l))
        aligned_feat = torch.stack(aligned_feat, dim=1)  # (b, t, c, h, w)

        if not self.with_tsa:
            aligned_feat = aligned_feat.view(b, -1, h, w)
        feat = self.fusion(aligned_feat)

        out = self.reconstruction(feat)
        out = self.lrelu(self.pixel_shuffle(self.upconv1(out)))
        out = self.lrelu(self.pixel_shuffle(self.upconv2(out)))
        out = self.lrelu(self.conv_hr(out))
        out = self.conv_last(out)
        if self.hr_in:
            base = x_center
        else:
            base = F.interpolate(
                x_center, scale_factor=4, mode='bilinear', align_corners=False)
        out += base
        return out