【pytorch torchvision源码解读系列—5】DenseNet

pytorch框架中有一个非常重要且好用的包：torchvision，顾名思义这个包主要是关于计算机视觉cv的。这个包主要由3个子包组成，分别是：torchvision.datasets、torchvision.models、torchvision.transforms。

具体介绍可以参考官网：https://pytorch.org/docs/master/torchvision

具体代码可以参考github：https://github.com/pytorch/vision

torchvision.models这个包中包含alexnet、densenet、inception、resnet、squeezenet、vgg等常用经典的网络结构，并且提供了预训练模型，可以通过简单调用来读取网络结构和预训练模型。

今天我们来解读一下DenseNet的源码实现。如果对DenseNet不是很了解 可以查看这里的论文笔记

https://blog.csdn.net/sinat_33487968/article/details/83684453

DenseNet由Dense Block组成，而Dense Block石油DenseLayer组成，下面就是DenseLayer类，可以看到每一层卷积之前都是用了batchnorm和relu，然后就是1*1的卷积和3*3的卷积。bn_size参数是bottleneck layer瓶颈层数的乘法因子。bottleneck layer就是指在3*3的卷积之前的那层1*1的卷积，它降低输入的feature map的数量。drop_rate是drop out的概率大小。

注意到最后返回值是torch.cat([x, new_features], 1)，也就是包括自己本身和提取的feature层堆叠在一起。

class _DenseLayer(nn.Sequential):
    def __init__(self, num_input_features, growth_rate, bn_size, drop_rate):
        super(_DenseLayer, self).__init__()
        self.add_module('norm1', nn.BatchNorm2d(num_input_features)),
        self.add_module('relu1', nn.ReLU(inplace=True)),
        self.add_module('conv1', nn.Conv2d(num_input_features, bn_size *
                        growth_rate, kernel_size=1, stride=1, bias=False)),
        self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)),
        self.add_module('relu2', nn.ReLU(inplace=True)),
        self.add_module('conv2', nn.Conv2d(bn_size * growth_rate, growth_rate,
                        kernel_size=3, stride=1, padding=1, bias=False)),
        self.drop_rate = drop_rate

    def forward(self, x):
        new_features = super(_DenseLayer, self).forward(x)
        if self.drop_rate > 0:
            new_features = F.dropout(new_features, p=self.drop_rate, training=self.training)
        return torch.cat([x, new_features], 1)

有了DenseLayer，就能用它来构建Dense Block。解释一下一些参数的含义：growth_rate (int) - 代表每一层有多少个新的feature层添加进去 (paper里面是k)，num_input features（int） - 输入的feature层数，num_layers（int）- 代表了Dense Block有多少层DenseLayers。

for循环就是根据num_layers逐层添加Dense Block注意每一层的输入num_input features不一样，是因为每一个DenseLayer之后feature的数目都会增加growth rate。直观一点理解，其实就是每一层的DenseLayer附带了自己学习到的feature还有自己本身一起传入下一个DenseLayer。所以后面的输出的结果会不断的增加，下图不能体现数据流动这一点，因为下图不是数据本身，而是代表卷积层。

class _DenseBlock(nn.Sequential):
    def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate):
        super(_DenseBlock, self).__init__()
        for i in range(num_layers):
            layer = _DenseLayer(num_input_features + i * growth_rate, growth_rate, bn_size, drop_rate)
            self.add_module('denselayer%d' % (i + 1), layer)

DenseBlock的结构可以看下图。

在Dense Block与Dense Block之间，论文添加了Transition层，这个层的作用是为了进一步提高模型的紧凑性，可以减少过渡层的特征图数量。这里直接就把num_input_features压缩成num_output_features，使用一层卷积实现的。

class _Transition(nn.Sequential):
    def __init__(self, num_input_features, num_output_features):
        super(_Transition, self).__init__()
        self.add_module('norm', nn.BatchNorm2d(num_input_features))
        self.add_module('relu', nn.ReLU(inplace=True))
        self.add_module('conv', nn.Conv2d(num_input_features, num_output_features,
                                          kernel_size=1, stride=1, bias=False))
        self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2))