【深度学习基础】CSPNet——PyTorch实现CSPDenseNet和CSPResNeXt

1 论文关键信息

论文链接: CSPNet: A New Backbone that can Enhance Learning Capability of CNN

1.1 CSP结构

论文提出Cross Stage Partial（CSP）结构，其初衷是减少计算量并且增强梯度的表现。主要思想是：在输入block之前，将输入分为两个部分，其中一部分通过block进行计算，另一部分直接通过一个shortcut进行concatenate。

作者基于Dense block（对DenseNet结构不了解的同学可以查看我以前的博客：【深度学习基础】pytorch实现DenseNet亲身实践）讲述了CSP的结构，如下图：

图片来自于论文

上图中，上者为原始的Dense Block，下者为Partial Dense Block。假设将输入按照part_ratio=0.5的比例分成两部分，并且假设一个Dense Block的输入为whc，growth rate为k，block layers为m，则：对于原始的DenseBlock，其CIO为（Convolutional input/Output）为(cm) + ((mm+m)d)/2；而对于使用了CSP结构的结构来说，其CIO下降为(cm) + ((m*m+m)*d)/2。

作者在论文阐述了CSP结构的优点:
（1）加强CNN的学习能力；（2）减少计算瓶颈，现在的网络大多计算代价昂贵，不利于工业的落地；（3）减少内存消耗。

1.2 关于Partial Transition Layer

关于Patial Transition Layer，论文提出了三种不同的结构，如下图：

图片来自论文

Transition layer的含义同DenseNet，是一组1x1的卷积层。上图中，不同类型的transition layer决定了梯度的结构方式，并且各有其优势：（c）图Fusion First的先将两个部分进行concatenate，然后再进行输入到Transion layer中，采用这种做法会是的大量特梯度信息被重用，有利于网络学习；（d）图中Fusion Last先将部分特征输入Transition layer，然后再进行concatenate，这样做损失了部分的梯度重用，但是由于Transition的输入维度比（c）图少，大大减少计算复杂度。论文中CSPNet采用的是图（b）中的结构，其结合了（c）（d）的结合，进一步提升了学习能力，但是也进一步提高了一些计算复杂度。 作者在论文中给出其使用不同Partial Transition Layer的实验结果，如下图。具体使用哪种结构可以根据条件和使用场景进行调整。

图片来自论文

2 pytorch实现CSP-DenseNet，CSP-ResNeXt

2.1 CSP-DenseNet

CSPNet沿用网络的Architecture，只是对于基础block进行改造，根据之前博客介绍的DenseNet结构修改相应代码即可实现。

2.1.1 Partial Dense block

沿用之前介绍DenseNet写的Dense Block代码：

class DenseBlock(nn.Module):

    def __init__(self, input_channels, num_layers, growth_rate):
        super(DenseBlock, self).__init__()
        self.num_layers = num_layers
        self.k0 = input_channels
        self.k = growth_rate
        self.layers = self.__make_layers()

    def __make_layers(self):
        layer_list = []
        for i in range(self.num_layers):
            layer_list.append(nn.Sequential(
                BN_Conv2d(self.k0 + i * self.k, 4 * self.k, 1, 1, 0),
                BN_Conv2d(4 * self.k, self.k, 3, 1, 1)
            ))
        return layer_list

    def forward(self, x):
        feature = self.layers[0](x)
        out = torch.cat((x, feature), 1)
        for i in range(1, len(self.layers)):
            feature = self.layers[i](out)
            out = torch.cat((feature, out), 1)
        return out

Partial Dense Block的实现：

class CSP_DenseBlock(nn.Module):

    def __init__(self, in_channels, num_layers, k, part_ratio=0.5):
        super(CSP_DenseBlock, self).__init__()
        self.part1_chnls = int(in_channels * part_ratio)
        self.part2_chnls = in_channels - self.part1_chnls
        self.dense = DenseBlock(self.part2_chnls, num_layers, k)
        # trans_chnls = self.part2_chnls + k * num_layers
        # self.transtion = BN_Conv2d(trans_chnls, trans_chnls, 1, 1, 0)

    def forward(self, x):
        part1 = x[:, :self.part1_chnls, :, :]
        part2 = x[:, self.part1_chnls:, :, :]
        part2 = self.dense(part2)
        # part2 = self.transtion(part2)
        out = torch.cat((part1, part2), 1)
        return out

上面的代码采用的是Fusion Last的Partial transition layer，取消其中的三行注释就是论文中采用的Partial transition layer。由于，我之前的Transition layer是写在整体网络结构中，没有放在block中，所以就不实现Fusion first的形式了。

2.1.2 网络整体结构

在之前原始的DenseNet结构的基础上进行修改，增加一个partion-ratio参数来控制使用的block结构和通道划分的比例就可以了，代码：

class DenseNet(nn.Module):

    def __init__(self, layers: object, k, theta, num_classes, part_ratio=0) -> object:
        super(DenseNet, self).__init__()
        # params
        self.layers = layers
        self.k = k
        self.theta = theta
        self.Block = DenseBlock if part_ratio == 0 else CSP_DenseBlock   # 通过part_tatio参数控制block type
        # layers
        self.conv = BN_Conv2d(3, 2 * k, 7, 2, 3)
        self.blocks, patches = self.__make_blocks(2 * k)
        self.fc = nn.Linear(patches, num_classes)

    def __make_transition(self, in_chls):
        out_chls = int(self.theta * in_chls)
        return nn.Sequential(
            BN_Conv2d(in_chls, out_chls, 1, 1, 0),
            nn.AvgPool2d(2)
        ), out_chls

    def __make_blocks(self, k0):
        """
        make block-transition structures
        :param k0:
        :return:
        """
        layers_list = []
        patches = 0
        for i in range(len(self.layers)):
            layers_list.append(self.Block(k0, self.layers[i], self.k))
            patches = k0 + self.layers[i] * self.k  # output feature patches from Dense Block
            if i != len(self.layers) - 1:
                transition, k0 = self.__make_transition(patches)
                layers_list.append(transition)
        return nn.Sequential(*layers_list), patches

    def forward(self, x):
        out = self.conv(x)
        out = F.max_pool2d(out, 3, 2, 1)
        # print(out.shape)
        out = self.blocks(out)
        # print(out.shape)
        out = F.avg_pool2d(out, 7)
        # print(out.shape)
        out = out.view(out.size(0), -1)
        out = F.softmax(self.fc(out))
        return out


def csp_densenet_121(num_classes=1000):
    return DenseNet([6, 12, 24, 16], k=32, theta=0.5, num_classes=num_classes, part_ratio=0.5)


def csp_densenet_169(num_classes=1000):
    return DenseNet([6, 12, 32, 32], k=32, theta=0.5, num_classes=num_classes, part_ratio=0.5)


def csp_densenet_201(num_classes=1000):
    return DenseNet([6, 12, 48, 32], k=32, theta=0.5, num_classes=num_classes, part_ratio=0.5)


def csp_densenet_264(num_classes=1000):
    return DenseNet([6, 12, 64, 48], k=32, theta=0.5, num_classes=num_classes, part_ratio=0.5)

2.2 CSP-ResNeXt

论文也给出了CSP-ResNet结构。我们知道在ResNeXt的实现中，若groups设置为1，就和ResNet结构相似，这里仅搭建CSP-ResNeXt。

图片来自论文

讨论：论文中提到由于输入的通道减半了，CSPNet不需要使用bottleneck结构了。但是我查看DarkNet框架下的网络参数，发现还是1x1,3x3,1x1的组合结构。有点纳闷，不知道我是不是对bottleneck结构有误解。可能作者指的是bottleneck指的是输出通道比输入通道数量要少的卷积层。

作者给的code链接是darknet框架下的cfg参数，由于我不熟悉这个框架，有点看不清楚。这里，我是按照论文的思想和自己理解搭建CSP-ResNeXt，网络输入尺寸和整体结构仍沿用ResNeXt的,不是论文实验确实采用的确切网络结构，谨供参考和了解思路。

我们再回顾一下ResNeXt的整体结构，之前博客的传送门：

图片来自ResNeXt论文

ResNeXt的主题是一些不同尺寸的block组成的，共4组，以ResNeXt-53为例，每组blocks的数量分别为[3, 4, 6, 3]。根据CSPNet的思想去除bottleneck结构，采用Fusion First和Fusion Last的模式。针对每一组，可以设计如下的结构：

图片来源：自己画的~

由上图，之前ResNeXt中Residual block中的最后一层1x1滤波器数量进行了2倍的expansion，这里改成三层均是一样的数量，那么我们根据这个结构进行编码。

再次提醒，这个复现不是SoAT算法所用的确切CSPResNet。但是其block结构符合上图中CSPNet的基本思想，整体架构沿用ResNet架构——虽然目前不准备实验验证，但我想它应该具有一定精度。

2.2.0 BN_CONV_LeakyReLU

参考作者，激活函数采用了LeakyReLU，但在一些目标检测框架中会用Mish作为激活函数，以更高的计算代价获取网络更好的收敛性。

class BN_Conv2d_Leaky(nn.Module):
    """
    BN_CONV_LeakyRELU
    """

    def __init__(self, in_channels: object, out_channels: object, kernel_size: object, stride: object, padding: object,
                 dilation=1, groups=1, bias=False) -> object:
        super(BN_Conv2d_Leaky, self).__init__()
        self.seq = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride,
                      padding=padding, dilation=dilation, groups=groups, bias=bias),
            nn.BatchNorm2d(out_channels)
        )

    def forward(self, x):
        return F.leaky_relu(self.seq(x))

2.2.1 Residual block

在这里，将最后一层1x1的数量改为和之前两层一致。在shortcut的实现上，我用均值池化来应对stride不为1的情况，因为感觉我之前在ResNeXt写的在1x1卷积的时候设置stride处理方式会损失一定的输入信息。 代码如下：

class ResidualBlock(nn.Module):
    """
    Residual block for CSP-ResNeXt
    """
    def __init__(self, in_channels, cardinality, group_width, stride=1):
        super(ResidualBlock, self).__init__()
        self.out_channels = cardinality * group_width
        self.conv1 = BN_Conv2d_Leaky(in_channels, self.out_channels, 1, 1, 0)
        self.conv2 = BN_Conv2d_Leaky(self.out_channels, self.out_channels, 3, stride, 1, groups=cardinality)
        self.conv3 = nn.Conv2d(self.out_channels, self.out_channels, 1, 1, 0)
        self.bn = nn.BatchNorm2d(self.out_channels)

        # make shortcut
        layers = []
        if in_channels != self.out_channels:
            layers.append(nn.Conv2d(in_channels, self.out_channels, 1, 1, 0))
            layers.append(nn.BatchNorm2d(self.out_channels))
        if stride != 1:
            layers.append(nn.AvgPool2d(stride))
        self.shortcut = nn.Sequential(*layers)

    def forward(self, x):
        out = self.conv3(self.conv2(self.conv1(x)))
        out = self.bn(out)
        out += self.shortcut(x)
        return F.leaky_relu(out)

2.2.3 网络Architecture

1、把上图中设计的结构作为一个整体来实现，称之为Stem，代码如下：

class Stem(nn.Module):
    def __init__(self, in_channels, num_blocks, cardinality, group_with, stride=2):
        super(Stem, self).__init__()
        self.c0 = in_channels // 2
        self.c1 = in_channels - in_channels // 2
        self.hidden_channels = cardinality * group_with
        self.out_channels = self.hidden_channels * 2
        self.trans_part0 = nn.Sequential(BN_Conv2d_Leaky(self.c0, self.hidden_channels, 1, 1, 0), nn.AvgPool2d(stride))
        self.block = self.__make_block(num_blocks, self.c1, cardinality, group_with, stride)
        self.trans_part1 = BN_Conv2d_Leaky(self.hidden_channels, self.hidden_channels, 1, 1, 0)
        self.trans = BN_Conv2d_Leaky(self.out_channels, self.out_channels, 1, 1, 0)

    def __make_block(self, num_blocks, in_channels, cardinality, group_with, stride):
        strides = [stride] + [1] * (num_blocks-1)
        channels = [in_channels] + [self.hidden_channels] * (num_blocks-1)
        return nn.Sequential(*[ResidualBlock(c, cardinality, group_with, s)
                               for c, s in zip(channels, strides)])

    def forward(self, x):
        x0 = x[:, :self.c0, :, :]
        x1 = x[:, self.c0:, :, :]
        out0 = self.trans_part0(x0)
        out1 = self.trans_part1(self.block(x1))
        out = torch.cat((out0, out1), 1)
        return self.trans(out)

2、以下代码提供了CSP-ResNeXt的架构以及CSP-ResNeXt-53(32x4d)的接口：

class CSP_ResNeXt(nn.Module):
    def __init__(self, num_blocks, cadinality, group_width, num_classes):
        super(CSP_ResNeXt, self).__init__()
        self.conv0 = BN_Conv2d_Leaky(3, 64, 7, 2, 3)
        self.pool1 = nn.MaxPool2d(3, 2, 1)
        self.conv1 = BN_Conv2d_Leaky(64, 128, 1, 1, 0)
        self.stem0 = Stem(cadinality*group_width*2, num_blocks[0], cadinality, group_width, stride=1)
        self.stem1 = Stem(cadinality*group_width*4, num_blocks[1], cadinality, group_width*2)
        self.stem2 = Stem(cadinality*group_width*8, num_blocks[2], cadinality, group_width*4)
        self.stem3 = Stem(cadinality*group_width*16, num_blocks[3], cadinality, group_width*8)
        self.global_pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(cadinality*group_width*16, num_classes)

    def forward(self, x):
        out = self.conv0(x)
        out = self.pool1(out)
        out = self.conv1(out)
        out = self.stem0(out)
        out = self.stem1(out)
        out = self.stem2(out)
        out = self.stem3(out)
        out = self.global_pool(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return F.softmax(out)


def csp_resnext_50_32x4d(num_classes=1000):
    return CSP_ResNeXt([3, 4, 6, 3], 32, 4, num_classes)

3、搭建完毕，测试网络：

net = csp_resnext_50_32x4d()
summary(net, (3, 256, 256))

测试结果如下：