IndexNet

Indexed Network

​ 卷积网络中上采样算子可以用文章中提出的索引函数（Index function）统一，说明了插值上采样与反池化操作等算子都是索引函数的特殊情况，由于在图像任务中，不同的部分图像对操作的要求不一样（比如图像边界需要提出边界信息，插值的方法就不如反最大值池化的操作；图像内部我们又需要邻域像素点的信息进行上采样的操作，这种情况下反最大值池化又不如插值方法），所以我们不能在图像中简单地用一个算子进行上采样，文章中提出了用Index function来引导采样的操作，其中Index function是一个可学习的函数（通过图像的信息学习—feature map），这样就可以根据图像的内容信息和位置信息进行不同的采样操作。

index的定义与理解

​ 就是一个加权（对于图像内容的感知强度），用于进行矩阵乘法，表示对每个像素点的敏感程度。本文是由在 maxunpooling 的传递 index 的基础上提出的，所以本文作者才会称其为 indexNet ，在文章中的index都是"soft index"，对 index 中每个元素: $i$，$i \in [0,1]$

对于每一个feature map $\Large \mathscr{x}$，IndexNet会产生两个Index function分别用于上采样与下采样：其实就是产生两个卷积核进行采样：

• 下采样操作 $\mathcal{IP}$ (indexed pooling)：设$I(x)$是 $x$ 通过 IndexNet 产生的索引（由下采样前的 feature map 产生，产生的 index map 大小与下采样前的 feature map 大小相同，因此在 $\mathcal{IP}$ 操作之后要进行一次平均值池化满足下采样后的大小要求）$E \in \mathbb{R}^{k \times k}$ 是下采样像素点的方形邻域（感受野），那么对$E$ 进行下采样，得到的结果为$\mathcal{IP}(E)=\sum_{x \in E}I(x)x$ ，对于每个区域 $E$，$I(x)$ 都会经过一个 softmax 保证图像强度连续性，从这里就可以看出最大值池化和平均值池化都是IndexNet的特例了。注：这样产生的 feature map 需要乘一个常数，在下采样率为2的时候整体 feature map 乘4去补偿平均值采样

• 上采样操作 $\mathcal{IU}$ (indexed upsampling)：它是下采样操作的反算子，它先将上采样前低分辨率的 feature map 进行 NN 插值到和上采样后的大小，以匹配产生的 $I(d)$ 的大小（这一步就相对于CARAFE不同了，它考虑了spatial information，因为信息来自于encoder）即得上采样区域，即 $\mathcal{IU}(d)=I(d) \otimes E$，$E$ 的定义同上面的下采样。

文章中还根据训练的参数量将Index function分类：

• 每个层通道是否共享一个Index function将Index Networks分为HINS(holistic index networks) 和 DINS(depthwise index networks)，对于原始$H \times W \times C$大小的feature map，HINS会将feature map特征提取输出为$H^{'} \times W^{'} \times 1$大小的Index block，而DINS会将feature map特征提取输出为$H^{'} \times W^{'} \times C$的Index block
• 再对DINS进行分类，由于输入feature map与输出Index block的通道数都是$C$，因此我们考虑使用三维卷积还是二维卷积，即输出Index block的对通道$C$的每一个切片是否考虑与feature map对通道$C$的每一个切片一一对应，还是考虑feature map的全部层，于是将DINS分类为O2O DINS(one to one)和M2O DINS(many to one)
• 我们再对O2O DINS进行分类，考虑在上采样与下采样操作中，我们是否采用同一个参数的Index block，于是将O2O DINS进行分类为Modelwise O2O DINS和Stagewise O2O DINS
• 再对Stagewise O2O DINS进行分类，如果我们在不同阶段的采样过程中（不只是上下采样）都需要不同的Index block（即是否共享参数），我们再将Stagewise O2O DINS分类为Shared Stagewise O2O DINS和unshared Stagewise O2O DINS

​ 其中M2O的效果最好，但是同时参数量也最大；HINs比DINs更加灵活，对于网络的设计来说自由度更高。

IndexNet与CARAFE的比较

相同点：

• 它们都是基于内容生成的动态的，具有可学习参数的采样核。

• CARAFE是IndexNet的特殊形式：

不同点

• 下采样时，IndexNet产生的Index function是基于前面高分辨率的feature map产生的，这样就不会丢失空间位置信息。

• CARAFE基于下采样后的低分辨率图产生，因此具有丰富的语义信息

• IndexNet的下采样与上采样方式是相互作用的，而CARAFE上下采样方式毫无关系，这一点可能可以引导以后的研究！

import torch
import torch.nn as nn
import torch.nn.functional as F
import einops

class HolisticIndexBlock(nn.Module):
    def __init__(self, inp, use_nonlinear=False, use_context=False,
                 use_batch_norm=False, ratio=2):
        super(HolisticIndexBlock, self).__init__()
        if use_context:
            kernel_size, padding = 4, 1
        else:
            kernel_size, padding = 2, 0

        c_mid = inp * ratio
        c_out = ratio**2
        self.ratio = ratio

        if use_batch_norm:
            BatchNorm2d = nn.BatchNorm2d
        else:
            BatchNorm2d = None

        if use_nonlinear:
            self.indexnet = nn.Sequential(
                    nn.Conv2d(inp, c_mid, kernel_size=kernel_size, stride=2,
                              padding=padding, bias=False),
                    BatchNorm2d(c_mid),
                    nn.ReLU6(inplace=True),
                    nn.Conv2d(c_mid, c_out, kernel_size=1, padding=0, bias=False)
                )
        else:
            self.indexnet = nn.Conv2d(inp, c_out, kernel_size=kernel_size,
                                      stride=2, padding=padding, bias=False)


    def forward(self, x):
        x = self.indexnet(x)
        y = F.sigmoid(x)
        z = F.softmax(y ,dim=1)
        # 可以直接使用stride=1的卷积，这样后续就不需要将通道维数补偿到图像大小上去
        idx_encoder = einops.rearrange(z, 'b (r1 r2) h w -> b 1 (h r1) (w r2)', r1=self.ratio)
        idx_decoder = einops.rearrange(y, 'b (r1 r2) h w -> b 1 (h r1) (w r2)', r1=self.ratio)
        return idx_encoder, idx_decoder


class DepthwizeIndexBlock(nn.Module):
    def __init__(self, inp, use_nonlinear=False, use_context=False,
                 use_batch_norm=False, ratio=2, mode='O2O'):
        super(DepthwizeIndexBlock, self).__init__()
        
        if 'M2O' in mode:
            grp = 1
        elif 'O2O' in mode:
            grp = inp
        else:
            assert "MODE ERROR"
        self.ratio = ratio
        self.grp = grp

        self.indexnets = nn.ModuleList([
            self._build_index_block(inp, use_nonlinear, use_context, use_batch_norm)
            for _ in range(ratio**2)
        ])
    
    def _build_index_block(self, inp, use_nonlinear, use_context, use_batch_norm):
        if use_context:
            kernel_size, padding = 4, 1
        else:
            kernel_size, padding = 2, 0
        
        if use_batch_norm:
            BatchNorm2d = nn.BatchNorm2d
        else:
            BatchNorm2d = None
        
        if use_nonlinear:
            return nn.Sequential(
                # 输入通道被分为grp组，每组之间共享一组卷积核
                # 其中每个组的通道只与同组的卷积核进行卷积操作，而不与其他组的卷积核发生作用
                nn.Conv2d(inp, inp, kernel_size=kernel_size, stride=2, 
                          padding=padding, groups=self.grp),
                BatchNorm2d(inp),
                nn.ReLU6(inplace=True),
                nn.Conv2d(inp, inp, kernel_size=1, stride=1, padding=0, groups=self.grp)
            )
        else:
            return nn.Conv2d(inp, inp, kernel_size=kernel_size, padding=padding,
                             stride=2, groups=self.grp)
    
    def forward(self, x):
        x_list = [indexnet(x) for indexnet in self.indexnets]
        x = torch.cat(x_list, dim=1)

        y = F.sigmoid(x)
        z = F.softmax(x, dim=2)

        idx_en = einops.rearrange(y, 'b (r1 r2 c) h w -> b c (h r1) (w r2)',
                             r1=self.ratio, r2=self.ratio)
        idx_de = einops.rearrange(z, 'b (r1 r2 c) h w -> b c (h r1) (w r2)',
                             r1=self.ratio, r2=self.ratio)
        return idx_en, idx_de

if __name__ == '__main__':
    DIN = DepthwizeIndexBlock(16, mode='M2O')
    x = torch.Tensor(3, 16, 24, 24)
    y, z = DIN(x)
    
    x1 = torch.einsum('bcij,bcij->bcij', [x,y])
    # x = F.interpolate(x, scale_factor=0.5, mode='area')
    x1 = F.avg_pool2d(x1, kernel_size=2, stride=2)
    x1 = x1*4

    x2 = F.interpolate(x1, scale_factor=2, mode='nearest')
    x2 = torch.einsum('bcij,bcij->bcij', [x2,z])
    print(x1.size())
    print(x2.size())

代码经验：

• 对于一些不确定是否存在的层，将这些层设置为None就可以跳过他们，而不需要写两个nn.Sequence
• 对于重复性的层，合理使用列表生成式写法减少代码量

参考文献：

Index Networks