https://arxiv.org/pdf/1905.02188 特征上采样是许多现代卷积网络架构(例如特征金字塔)中的关键操作。其设计对于诸如目标检测和语义/实例分割等密集预测任务至关重要。在本文中,我们提出了内容感知特征重组(CARAFE),这是一种通用、轻量级且高度有效的算子,以实现这一目标。CARAFE具有几个吸引人的特性:(1)大视野。与仅利用子像素邻域的前期工作(例如双线性插值)不同,CARAFE可以在大感受野内聚合上下文信息。(2)内容感知处理。与对所有样本使用固定核(例如反卷积)不同,CARAFE支持针对实例的内容感知处理,可以即时生成自适应核。(3)轻量级且计算速度快。CARAFE引入的计算开销很小,并且可以轻松集成到现代网络架构中。我们在目标检测、实例/语义分割和图像修复的标准基准上进行了综合评估。CARAFE在所有任务中都表现出一致且显著的增益(分别为、、、),且计算开销可忽略不计。它有望成为未来研究的强大构建块。代码和模型可在https://github.com/open-mmlab/mmdetection获取。
特征上采样是深度神经网络中最基本的操作之一。一方面,对于密集预测任务(例如超分辨率[7,20]、图像修复[13,32]和语义分割[43,5])中的解码器,高级/低分辨率特征图会上采样以匹配高分辨率监督。另一方面,特征上采样还涉及将高级/低分辨率特征图与低级/高分辨率特征图进行融合,这在许多最先进架构中被广泛采用,例如特征金字塔网络[21]、U-Net[34]和堆叠沙漏网络[29]。因此,设计有效的特征上采样算子成为一个关键问题。
最常用的特征上采样算子是最近邻和双线性插值,它们采用像素之间的空间距离来指导上采样过程。然而,最近邻和双线性插值仅考虑子像素邻域,无法捕获密集预测任务所需的丰富语义信息。另一种实现自适应上采样的方法是反卷积[30]。反卷积层作为卷积层的逆算子,学习一组与实例无关的上采样核。但是,它有两个主要缺点。首先,反卷积算子在整个图像上应用相同的核,而不考虑底层内容。这限制了其响应局部变化的能力。其次,当使用大核大小时,它带有大量参数和繁重的计算工作量。这使得其难以覆盖超出小邻域的更大区域,从而限制了其表达能力和性能。
在本文中,我们超越了这些限制,并寻求一种特征上采样算子,其能够1)在大感受野内聚合信息,2)即时适应特定实例的内容,以及3)保持计算效率。为此,我们提出了一种轻量级且高度有效的算子,称为内容感知特征重组(CARAFE)。具体而言,CARAFE通过加权组合,在每个位置为中心的预定区域内重新组合特征,其中权重以内容感知的方式生成。此外,对于每个位置,有多组这样的上采样权重。特征上采样通过将生成的特征重新排列为空间块来完成。
请注意,这些空间自适应权重不是作为网络参数学习的。相反,它们是使用具有softmax激活的轻量级全卷积模块即时预测的。图1揭示了CARAFE的工作机制。通过CARAFE上采样后,特征图可以更准确地表示对象的形状,从而使模型能够预测更好的实例分割结果。我们的CARAFE不仅在空间上对特征图进行上采样,还学习增强其判别能力。
为了证明CARAFE的普遍有效性,我们使用主流架构在广泛的密集预测任务上进行了全面评估,即目标检测、实例分割、语义分割和图像修复。在MS COCO 2018 test-dev数据集上,CARAFE可以使Faster RCNN[33]在目标检测上的性能提升,使Mask RCNN[9]在实例分割上的性能提升。在ADE20k[47, 48]验证集上的语义分割任务中,CARAFE使UperNet[38]的性能提升了,在Places[46]验证集上的图像修复任务中,使Global&Local[13]的PSNR提升了1.1 dB。当将具有256个通道的特征图上采样两倍时,CARAFE引入的计算开销仅为 FLOPs,相比之下,反卷积的计算开销为 FLOPs。在所有任务上取得的显著增益表明,CARAFE是一种有效且高效的特征上采样算子,未来有很大的潜力成为强有力的研究基石。
上采样算子。最常用的上采样方法是最近邻插值和双线性插值。这些插值方法利用距离来衡量像素之间的相关性,并使用手工设计的上采样核。在深度学习时代,提出了几种使用可学习算子对上采样特征图的方法。例如,反卷积[30]是卷积的逆算子,在这些可学习上采样器中最为著名。Pixel Shuffle[35]提出了一种不同的上采样器,它将通道空间上的深度重塑为空间上的宽度和高度。最近,[26]提出了引导上采样(GUM),通过采样具有可学习偏移量的像素来进行插值。然而,这些方法要么仅利用了小邻域内的上下文信息,要么需要昂贵的计算来进行自适应插值。在超分辨率和去噪领域,其他工作[27,16,11]也探索了在低级视觉任务中空间上使用可学习核。本着类似的设计精神,本文展示了内容感知特征重组在几个视觉感知任务中上采样的有效性和工作机制,并提供了一种轻量级解决方案。
密集预测任务。目标检测是定位带有边界框的物体的任务,实例分割进一步要求预测实例级掩码。Faster-RCNN[33]引入了区域提议网络(RPN)用于端到端训练,该网络通过引导锚定方案[37]得到了进一步改进。[21, 24, 17, 45, 31]利用多尺度特征金字塔来处理不同尺度的物体。通过添加额外的掩码预测分支,Mask-RCNN[9]及其变体[1, 12]获得了令人满意的像素级结果。语义分割[25, 19]要求对给定图像进行像素级语义预测。PSPNet[43]引入了多个网格尺度的空间池化,而UperNet[38]基于PSPNet设计了一个更通用的框架。图像或视频修复[42, 40, 39]是填充输入图像缺失区域的一个经典问题。U-net[34]在最近的工作中[13, 36]很受欢迎,并采用了多个上采样算子。Liu等人[23]引入了部分卷积层,以减轻缺失区域对卷积层的影响。我们的CARAFE在各种密集预测任务中展示了普遍的有效性。
特征上采样是许多现代卷积网络架构中的关键算子,这些架构是为包括目标检测、实例分割和场景解析在内的任务开发的。在这项工作中,我们提出了内容感知特征重组(CARAFE)来上采样特征图。在每个位置上,CARAFE可以利用底层内容信息来预测重组核,并在预定义的附近区域内重新组合特征。得益于内容信息,CARAFE可以在不同位置使用自适应和优化的重组核,并且比主流上采样算子(如插值或反卷积)实现更好的性能。
CARAFE作为具有内容感知核的重组算子,包含两个步骤。第一步是根据内容为每个目标位置预测一个重组核,第二步是使用预测的核重新组合特征。给定大小为的特征图和上采样比例(假设为整数),CARAFE将生成大小为的新特征图。对于输出上的任何目标位置,在输入上都有一个对应的源位置,其中。这里,我们将表示为以位置为中心、大小为的的子区域,即的邻域。
在第一步中,核预测模块根据的邻域为每个位置预测一个位置感知核,如公式(1)所示。重组步骤如公式(2)所示,其中是内容感知重组模块,它使用核重新组合的邻域:
在以下部分中,我们将详细说明和的细节。
核预测模块负责以内容感知的方式生成重组核。在上的每个源位置对应于上的个目标位置。每个目标位置需要一个的重组核,其中是重组核的大小。因此,该模块将输出大小为的重组核,其中。
核预测模块由三个子模块组成,即通道压缩器、内容编码器和核归一化器,如图2所示。通道压缩器减少了输入特征图的通道数。然后,内容编码器将压缩后的特征图作为输入,对内容进行编码以生成重组核。最后,核归一化器对每个重组核应用softmax函数。以下是对这三个子模块的详细解释。
通道压缩器。我们采用一个的卷积层来将输入特征通道从压缩到。减少输入特征图的通道数可以减少后续步骤中的参数数量和计算成本,使CARAFE更高效。在相同的预算下,也可以使用更大的内核大小为内容编码器。实验结果表明,在可接受的范围内减少特征通道不会损害性能。
内容编码器。我们使用大小为的卷积层,根据输入特征的内容生成重组核。编码器的参数为。直观上,增加可以扩大编码器的感受野,并利用更大区域内的上下文信息,这对于预测重组核很重要。然而,计算复杂度随着内核大小的平方增长,而来自更大内核大小的收益则不会。通过我们在5.3节的研究,经验公式在性能和效率之间取得了良好的平衡。
核归一化器。在应用于输入特征图之前,每个的重组核都会通过softmax函数在空间上进行归一化。归一化步骤强制内核值的和为1,这是在局部区域内的软选择。由于核归一化器,CARAFE不会执行任何重新缩放并改变特征图的平均值,这就是为什么我们提出的算子被称为特征的重组。
使用每个重组核,内容感知重组模块将通过函数重新组合局部区域内的特征。我们采用的简单形式,它只是一个加权和算子。对于目标位置和以为中心的相应正方形区域,重组如公式(3)所示,其中:
使用重组核,区域中的每个像素都会基于特征的内容而不是位置的距离,对上采样像素产生不同的贡献。由于局部区域内相关点的信息可以更多地被关注,因此重新组合后的特征图的语义可能比原始特征图更强。
在这里,我们讨论了CARAFE与动态滤波器[15]、空间注意力[3]、空间变换器[14]和可变形卷积[6]之间的关系,它们具有相似的设计理念但侧重点不同。
动态滤波器。动态滤波器根据网络的输入生成特定实例的卷积滤波器,然后将预测的滤波器应用于输入。动态滤波器和CARAFE都是内容感知算子,但它们之间的根本区别在于内核生成过程。具体来说,动态滤波器是一个两步卷积过程,其中额外的滤波器预测层和滤波层需要大量的计算。相反,CARAFE只是局部区域内特征的重新组合,而不学习跨通道的特征变换。假设输入特征图的通道数为,滤波器的内核大小为,则在动态滤波器中,每个位置预测的内核参数为。对于CARAFE,内核参数仅为。因此,它在内存和速度上更高效。
空间注意力。空间注意力预测一个与输入特征大小相同的注意力图,然后对每个位置的特征图进行重新缩放。我们的CARAFE通过加权和对局部区域内的特征进行重新组合。总之,空间注意力是一个具有逐点指导的重缩放算子,而CARAFE是一个具有局部区域指导的重新组合算子。空间注意力可以看作是CARAFE的一种特殊情况,其中重组核大小为1,不考虑核归一化器。
空间变换网络(STN)。STN根据输入特征图预测一个全局参数变换,并通过该变换对特征进行扭曲。然而,这个全局参数变换假设太强,无法表示复杂的空间变化;而且STN的训练很难。在这里,CARAFE使用位置特定的重组来处理空间关系,这能够实现更灵活的局部几何建模。
可变形卷积网络(DCN)。DCN也采用了学习几何变换的思想,并将其与常规卷积层相结合。它预测内核偏移量,而不是使用网格卷积内核。与动态滤波器类似,它也是一个重参数算子,计算成本比CARAFE高24倍。它也对参数初始化很敏感。
CARAFE可以无缝集成到需要上采样算子的现有框架中。在这里,我们介绍一些在主流密集预测任务中的应用。CARAFE通过引入可忽略的额外参数,在高层次和低层次任务(如目标检测、实例分割、语义分割和图像修复)中均提升了最先进方法的性能。
特征金字塔网络(FPN)是目标检测和实例分割领域的一种重要且有效的架构。它显著提高了Faster R-CNN和Mask R-CNN等流行框架的性能。FPN通过自顶向下的路径和横向连接构建具有强大语义的特征金字塔。在自顶向下的路径中,首先通过最近邻插值将低分辨率特征图上采样2倍,然后与高分辨率特征图进行融合,如图3所示。
我们提出在所有特征级别中用CARAFE替代最近邻插值。这一修改很平滑,无需额外更改。除了FPN结构外,Mask R-CNN在掩码头的末尾采用了一个反卷积层。它用于将预测的数字上采样到,以获得更精细的掩码预测。我们也可以使用CARAFE来替代反卷积层,从而进一步降低计算成本。
语义分割要求模型对整个图像输出逐像素级别的预测,因此通常更喜欢高分辨率特征图。在该任务中,上采样被广泛用于放大特征图并融合不同级别的语义信息。UperNet是语义分割的一个强大基线。它在以下三个组件中使用上采样,即PPM、FPN、FUSE。我们采用CARAFE替代它们原有的上采样器。
金字塔池化模块(PPM)。PPM是PSPNet中的关键组件,它分层地将输入特征图下采样到多个尺度,然后通过双线性插值将它们上采样回原始尺寸。最后,通过拼接将这些特征与原始特征进行融合。由于上采样比例非常大,我们采用两步策略,使用CARAFE在性能和效率之间取得平衡。首先,我们使用双线性插值将特征上采样到原始特征图尺寸的一半,然后使用CARAFE进一步将其上采样2倍。
特征金字塔网络(FPN)。与检测模型类似,UperNet也采用FPN来丰富特征语义。它只有四个不同的特征级别,步长分别为。我们以与4.1节相同的方式替换上采样算子。
多级特征融合(FUSE)。UperNet在FPN之后提出了一个多级特征融合模块。它通过双线性插值将P3、P4、P5上采样到与P2相同的尺寸,然后通过拼接融合这些不同级别的特征。该过程相当于一个顺序上采样-拼接过程,即首先将P5上采样到P4的尺寸并进行拼接,然后将拼接后的特征图上采样到P3的尺寸,以此类推。我们在这里用CARAFE替换顺序双线性上采样。
U-net架构在最近提出的图像修复方法中很受欢迎,如Global&Local [13]和Partial Conv [23]。在网络的后半部分有两个上采样算子。我们简单地将这两个上采样层替换为CARAFE并评估性能。对于Partial Conv,我们可以通过使用我们的内容感知重组核更新掩码来方便地保持CARAFE中的掩码传播。
数据集与评估指标。我们在几个重要的密集预测基准上评估了CARAFE。默认情况下,我们使用训练集进行训练,并在验证集上评估这些数据集的性能。 目标检测和实例分割。我们在具有挑战性的MS COCO 2017数据集上进行了实验。结果采用标准的COCO评估指标,即IoU从0.5到0.95的平均精度(mAP)。 语义分割。我们采用ADE20k基准来评估我们的方法在语义分割任务中的表现。结果通过平均IoU(mIoU)和像素准确率(P.A.)来衡量,分别表示预测与真实掩码之间的平均IoU和每像素分类准确率。 图像修复。我们采用Places数据集进行图像修复。使用L1误差(越低越好)和PSNR(越高越好)作为评估指标。 实现细节。除非另有说明,否则CARAFE在实验中采用一组固定的超参数,其中通道压缩器的为64,内容编码器的,。更多实现细节见补充材料。 目标检测和实例分割。我们在Faster RCNN和Mask RCNN上评估了CARAFE,主干网络为ResNet-50带FPN,并遵循Detectron [8]和MMDetection [2]的1x训练计划设置。语义分割。我们使用UperNet的官方实现,并采用相同的实验设置。 图像修复。我们采用Global&Local [13]和Partial Conv [23]作为基线方法来评估CARAFE。
目标检测与实例分割。我们首先通过将FPN中的最近邻插值替换为CARAFE(针对Faster RCNN和Mask RCNN),以及将Mask RCNN中掩码头的反卷积层替换为CARAFE来评估我们的方法。如表1所示,CARAFE使Faster RCNN的bbox AP提高了,使Mask RCNN的掩码AP提高了。、、的提升均超过,这表明它对各种目标尺度都有益。
如图1所示的定性结果支持了我们令人鼓舞的性能。我们可视化了FPN自上而下路径中的特征图,并将CARAFE与基线(即最近邻插值)进行了比较。显然,通过内容感知重组,特征图更具判别性,并且预测出了更精确的目标掩码。在图4中,我们展示了一些实例分割结果的示例,比较了基线和CARAFE。
为了研究不同上采样算子的有效性,我们在Faster RCNN中使用了不同的算子在FPN中进行上采样,并进行了大量实验。结果如表2所示。对于“N.C.”和“B.C.”,分别表示“最近邻+卷积”和“双线性+卷积”,我们在相应的上采样后添加了一个额外的卷积层。“Deconv”、“Pixel Shuffle”(表示为“P.S.”)、“GUM”是三种具有代表性的基于学习的上采样方法。我们还比较了这里的“空间注意力”(表示为“S.A.”)。CARAFE在这些上采样算子中取得了最高的AP,且FLOPs和参数相对较少,这证明了它既有效又高效。“最近邻+卷积”和“双线性+卷积”的结果表明,额外的参数并没有带来显著的增益。“Deconv”、“Pixel Shuffle”、“GUM”和“空间注意力”的性能均低于CARAFE,这表明设计有效的上采样算子至关重要。
除了FPN这种金字塔特征融合结构外,我们还探索了掩码头中的不同上采样算子。在典型的Mask R-CNN中,采用反卷积层将RoI特征上采样2倍。为了公平比较,我们没有对FPN进行任何修改,只将反卷积层替换不同的为算子。由于我们只修改了掩码预测分支,因此以掩码AP的形式报告性能,如表3所示。在这些方法中,CARAFE在实例分割方面取得了最佳性能。
在表4中,我们报告了在Mask RCNN的FPN和掩码头中分别采用CARAFE的目标检测和实例分割结果。这些实验均取得了一致的改进。 语义分割。我们将UperNet中的上采样器替换为CARAFE,并在ADE20k基准上评估结果。如表5所示,CARAFE将单尺度测试的mIoU从大幅提升至。值得注意的是,带有CARAFE的UperNet也取得了比最近的强大基线(如PSPNet [43]和PSANet [44])更好的性能。
我们按照4.2节所述,逐步研究了修改UperNet中不同组件的有效性。表6中的结果表明,CARAFE对所有三个组件都有帮助,且它们的组合带来了进一步的增益。 图像修复。我们证明了CARAFE在图像修复等低级任务中也是有效的。通过在两个强大的基线Global&Local [13]和Partial Conv [23]中将上采样算子替换为CARAFE,我们观察到这两种方法都有显著改进。如表7所示,我们的方法在PSNR指标上使两个基线分别提高了1.1 dB和0.2 dB。
模型设计与超参数。我们研究了模型设计中超参数的影响,即压缩通道、编码器核大小和重组核大小。我们还测试了核归一化器中不同的归一化方法。我们使用带有ResNet-50主干的Faster RCNN对设计和设置进行了消融研究,并在COCO 2017验证集上评估了结果。
为了实现高效设计,我们首先分析了由浮点运算数(FLOPs)衡量的计算复杂度。当以因子对输入通道为的特征图进行上采样时,CARAFE的每像素FLOPs计算为,参见[28]。
我们在通道压缩器中尝试了不同的值。此外,我们还尝试移除通道压缩器模块,这意味着内容编码器直接使用输入特征来预测重组核。表8中的实验结果表明,将压缩到64不会导致性能下降,反而更加高效。进一步减小会导致性能略有下降。在没有通道压缩器的情况下,它可以达到相同的性能,这证明了通道压缩器可以在不损害性能的情况下加快核预测速度。基于以上结果,我们默认将设置为64,以在性能和效率之间取得平衡。
然后,我们研究了和的影响。直观上,增加也需要更大的,因为内容编码器需要一个大的感受野来预测一个大的重组核。如表9所示,同时增加和可以提高性能,而只增加其中一个则不会。我们总结了一个经验公式,这是所有设置中的一个好选择。尽管采用更大的核大小是有帮助的,但我们默认设置和,以在性能和效率之间取得平衡。
除了softmax函数外,我们还测试了核归一化器中的其他替代方案,如sigmoid或带归一化的sigmoid。如表10所示,“Softmax”和“Sigmoid Normalized”具有相同的性能,并且优于“Sigmoid”,这表明将重组核归一化为和为1是至关重要的。
CARAFE的工作原理。我们进行了进一步的定性研究,以弄清楚CARAFE是如何工作的。我们使用训练好的采用CARAFE作为上采样算子的Mask RCNN模型,在图5中可视化了重组过程。在特征金字塔网络(FPN)结构中,低分辨率特征图将连续上采样几次以达到更高的分辨率,因此上采样特征图中的一个像素会从更大的区域重新组合信息。我们在高分辨率特征图中采样了一些像素,并查看它们是从哪些相邻像素重新组合而来的。绿色圆圈表示示例位置,红色点表示在重组过程中高度加权的源。从图中可以清楚地看出,CARAFE是内容感知的。它倾向于重新组合具有相似语义信息的点。人体上的一个位置更倾向于来自同一人体的其他点,而不是其他物体或附近的背景。对于具有较弱语义的背景区域中的位置,重组更加均匀或只是偏向于具有相似低级纹理特征的点。
我们提出了内容感知特征重组(CARAFE),这是一种通用、轻量级且高度有效的上采样算子。它在目标检测、实例/语义分割和图像修复的标准基准上分别一致地将性能提高了、、、。更重要的是,CARAFE引入的计算开销很小,并且可以轻松集成到现代网络架构中。未来的研究方向包括探索CARAFE在低级视觉任务(如图像恢复和超分辨率)中的适用性。
致谢。本工作部分得到了商汤科技集团(中大协议编号:TS1610626和TS1712093)、香港研究资助局(GRF)(编号:14236516和14203518)、新加坡教育部(MOE)AcRF一级(M4012082.020)、南洋理工大学(NTU)学生资助基金(SUG)和NTU NAP的联合研究资助。
目标检测和实例分割。我们在以ResNet-50为主干网络的Faster RCNN[33]和Mask RCNN[9]上评估CARAFE。这些方法使用了FPN[21]。在训练和推理中,我们调整输入图像的大小,使其短边有800像素或长边有1333像素,同时保持其纵横比不变。我们采用同步SGD,初始学习率为0.02,动量为0.9,权重衰减为0.0001。我们在8个GPU上使用16的批量大小(每个GPU 2张图像)。遵循Detectron[8]和MMDetection[2]的1倍训练计划,我们总共训练12个周期,并在第8个和第11个周期将学习率降低0.1倍。
语义分割。我们使用以ResNet-50为主干网络的UperNet的官方实现{ }^{3} [38]。在训练期间,输入图像的短边大小随机从{300,375,450, 525,600}中选择。在推理中,为了公平比较,我们采用单尺度测试,并将图像的短边设置为450像素。在训练和推理中,图像长边的最大长度均设置为1200。我们采用同步SGD,初始学习率为0.02,动量为0.9,权重衰减为0.0001。我们在8个GPU上使用16的批量大小(每个GPU 2张图像),并采用同步批量归一化,这是语义分割中的常见做法。遵循[4],我们采用“poly”学习率策略,其中当前迭代的学习率等于初始学习率乘以。我们将power设置为0.9,并总共训练20个周期。
图像修复。我们采用Global&Local[13]中的生成器和判别器网络作为基线。我们的生成器以带有掩码区域的图像作为输入,并输出缺失区域的预测。然后,我们通过将预测图像与输入图像结合。最后,将组合后的输出输入到判别器中。我们对基线模型进行了简单修改,以获得更好的生成质量。与原始模型使用两个判别器相比,我们仅在修复区域使用一个PatchGAN风格的判别器[18]。这种修改可以获得更好的图像质量。
为了公平比较并考虑实际应用,我们使用[41]引入的自由形式掩码作为二进制掩码。对于Partial Conv[23],我们只需在生成器中将卷积层替换为官方的Partial Conv模块。在训练期间,我们采用Adam求解器,学习率为0.0001,其中,。训练批量大小为32。输入和输出在范围内线性缩放。
在图6中,我们通过更多示例展示了CARAFE如何进行内容感知重组。在FPN结构的自上而下路径中,红色单元通过CARAFE重组为绿色中心单元。
附录C. 可视化结果比较 目标检测和实例分割。如图7所示,我们提供了Mask RCNN基线和带有CARAFE的Mask RCNN在COCO[22] 2017验证集上的更多目标检测和实例分割结果比较。
语义分割。我们在图8中比较了ADE20k[47]验证集上UperNet基线和带有CARAFE的UperNet的语义分割结果。
图像修复。在图9中,我们展示了Places[46]验证集上Global&Local基线和带有CARAFE的Global&Local的图像修复结果比较。
在这里插入图片描述
在这里插入图片描述
class CARAFE(nn.Module):
def __init__(self, c, k_enc=3, k_up=5, c_mid=64, scale=2):
""" The unofficial implementation of the CARAFE module.
The details are in "https://arxiv.org/abs/1905.02188".
Args:
c: The channel number of the input and the output.
c_mid: The channel number after compression.
scale: The expected upsample scale.
k_up: The size of the reassembly kernel.
k_enc: The kernel size of the encoder.
Returns:
X: The upsampled feature map.
"""
super(CARAFE, self).__init__()
self.scale = scale
self.comp = Conv(c, c_mid)
self.enc = Conv(c_mid, (scale*k_up)**2, k=k_enc, act=False)
self.pix_shf = nn.PixelShuffle(scale)
self.upsmp = nn.Upsample(scale_factor=scale, mode='nearest')
self.unfold = nn.Unfold(kernel_size=k_up, dilation=scale,
padding=k_up//2*scale)
def forward(self, X):
b, c, h, w = X.size()
h_, w_ = h * self.scale, w * self.scale
W = self.comp(X) # b * m * h * w
W = self.enc(W) # b * 100 * h * w
W = self.pix_shf(W) # b * 25 * h_ * w_
W = torch.softmax(W, dim=1) # b * 25 * h_ * w_
X = self.upsmp(X) # b * c * h_ * w_
X = self.unfold(X) # b * 25c * h_ * w_
X = X.view(b, c, -1, h_, w_) # b * 25 * c * h_ * w_
X = torch.einsum('bkhw,bckhw->bchw', [W, X]) # b * c * h_ * w_
return X