本篇文章对多尺度层次化的transformer有新的诠释,以往的多尺度层次化transformer主干往往每个阶段的特征图只参与本阶段的计算,在本文中,每个阶段的特征图都要随着主干推进而更新并且参与其他特征图的计算。
下图是整体架构图。
值得注意的点有这么几个,一是HRViT的transformer模块使用的是HRattention和MixCFN代替传统的self-attention与FFN。二是每个阶段中不同尺度之间的交互,就是上图的红绿线,这个在文中被称为Cross-resolution fusion layer。
这个fusion layer其实就是上图中红绿线,代表着不同尺度的特征图进行交互。比如说从stage3到stage4的交互相当于有三个输入块,要获得4个输出块。依次处理每个输出块。要将输入中所有块经过交互处理后叠加成为输出块。交互处理我们可以看出来有三种情况,一种是HxW较大的特征图试图融合HxW较小的特征图,即红线,需要较小的特征图进行上采样。另一种则是相反,还有一种是同规格的特征图。这三种情况怎么处理呢。一般是从小到大是通过1x1的卷积通道数缩小到目标输出模块的通道数,再用最近点插值到目标模块尺寸,从大到小则是先用depth-wise 卷积缩小尺寸,再用1x1conv调整通道数。而相同规格尺寸的输入模块不需要处理。最后,一个输出模块就是输入模块数量的经过处理后获得相同的尺寸,通道然后相加。
class HRViTFusionBlock(nn.Module):
def __init__(
self,
in_channels: Tuple[int] = (32, 64, 128, 256),
out_channels: Tuple[int] = (32, 64, 128, 256),
act_func: nn.Module = nn.GELU,
norm_cfg: Dict = dict(type="SyncBN", requires_grad=True),
with_cp: bool = False,
) -> None:
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.act_func = act_func
self.norm_cfg = norm_cfg
self.with_cp = with_cp
if with_cp:
momentum = 1 - 0.9 ** 0.5
self.norm_cfg.update(dict(momentum=momentum))
self.n_outputs = len(out_channels)
self._build_fuse_layers()
def _build_fuse_layers(self):
self.blocks = nn.ModuleList([])
n_inputs = len(self.in_channels)
for i, outc in enumerate(self.out_channels):
blocks = nn.ModuleList([])
start = 0
end = n_inputs
for j in range(start, end):
inc = self.in_channels[j]
if j == i:
blocks.append(nn.Identity())
elif j < i:
block = [
nn.Conv2d(
inc,inc,kernel_size=2 ** (i - j) + 1,stride=2 ** (i - j),dilation=1,
padding=2 ** (i - j) // 2,groups=inc,bias=False,),build_norm_layer(self.norm_cfg, inc)[1],
nn.Conv2d(
inc,outc,kernel_size=1,stride=1,dilation=1,padding=0,groups=1,bias=True,
),
build_norm_layer(self.norm_cfg, outc)[1],
]
blocks.append(nn.Sequential(*block))
else:
block = [
nn.Conv2d(
inc,outc,kernel_size=1,stride=1,dilation=1,
padding=0,groups=1,bias=True,
),build_norm_layer(self.norm_cfg, outc)[1],]
block.append(
nn.Upsample(scale_factor=2 ** (j - i),mode="nearest",),)
blocks.append(nn.Sequential(*block))
self.blocks.append(blocks)
self.act = nn.ModuleList([self.act_func() for _ in self.out_channels])
def forward(self,x: Tuple[Tensor,],) -> Tuple[Tensor, ]:
out = [None] * len(self.blocks)
n_inputs = len(x)
for i, (blocks, act) in enumerate(zip(self.blocks, self.act)):#output_num
start = 0
end = n_inputs
for j, block in zip(range(start, end), blocks):#input_num
out[i] = block(x[j]) if out[i] is None else out[i] + block(x[j])#不同尺度融合后相加
out[i] = act(out[i])
return out
。
attention块首先要把输入特征图按维度一分为二。一半送去做horizental attention,另一半送去做vertical attention.所谓horizental attention就是设置一个大小为wsW的条状窗口,把q,k,v的shape从N,C调整到wsW,C来做attention,另外取QKV的时候只取QK,共享K和V.其他就没什么了。attention也是正常的计算过程。在计算完成后,要和vertical attention拼接,这样维度又恢复了。
另外本文还提出了一种新的mix-cfn来取代FFN,rc代表hidden dimension。将rc一拆为二后分别做3x3和5x5DW卷积再拼接。代码与图有些细微的区别,图中是通过两次linear得到两个特征图,代码是先Linear得到rc特征图,再一拆为二。