DenseNet論文翻譯及pytorch實現解析（下）

02-05

前言：pytorch提供的DenseNet代碼是在ImageNet上的訓練網路。

根據前文所述，DenseNet主要有DenseBlock和Transition兩個模塊。

DenseBlock

實現代碼：

class _DenseLayer(nn.Sequential):#卷積塊：BN->ReLU->1x1Conv->BN->ReLU->3x3Convn def __init__(self, num_input_features, growth_rate, bn_size, drop_rate):n"num_input_features:輸入特徵圖個數ngrowth_rate:增長速率，第二個卷積層輸出特徵圖ngrow_rate * bn_size:第一個卷積層輸出特徵圖ndrop_rate:dropout失活率n"n super(_DenseLayer, self).__init__()n self.add_module(norm.1, nn.BatchNorm2d(num_input_features)),n self.add_module(relu.1, nn.ReLU(inplace=True)),n self.add_module(conv.1, nn.Conv2d(num_input_features, bn_size *n growth_rate, kernel_size=1, stride=1, bias=False)),n self.add_module(norm.2, nn.BatchNorm2d(bn_size * growth_rate)),n self.add_module(relu.2, nn.ReLU(inplace=True)),n self.add_module(conv.2, nn.Conv2d(bn_size * growth_rate, growth_rate,n kernel_size=3, stride=1, padding=1, bias=False)),n self.drop_rate = drop_ratenn def forward(self, x):n new_features = super(_DenseLayer, self).forward(x)n if self.drop_rate > 0:n new_features = F.dropout(new_features, p=self.drop_rate, training=self.training)n return torch.cat([x, new_features], 1)nnnclass _DenseBlock(nn.Sequential):n def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate):n"num_layers:每個block內dense layer層數"n super(_DenseBlock, self).__init__()n for i in range(num_layers):n layer = _DenseLayer(num_input_features + i * growth_rate, growth_rate, bn_size, drop_rate)n self.add_module(denselayer%d % (i + 1), layer)n

為便於理解，繪製了一個block（2個dense層）內信息的傳播過程：

如圖，不同的顏色表示信息的傳播過程。這就使得當前層的輸入是之前所有特徵圖（同一個block內）的拼接。

Transition

因為使用拼接的緣故，每經過一次拼接輸出通道數可能會激增。為了控制模型複雜度，這裡引入一個過渡塊。（引自李沐gluon中文教程）

實現代碼如下：

class _Transition(nn.Sequential):#過渡層，將特徵圖個數減半n def __init__(self, num_input_features, num_output_features):n"num_input_features:輸入特徵圖個數nnum_output_features:輸出特徵圖個數，為num_input_features//2n"n super(_Transition, self).__init__()n self.add_module(norm, nn.BatchNorm2d(num_input_features))n self.add_module(relu, nn.ReLU(inplace=True))n self.add_module(conv, nn.Conv2d(num_input_features, num_output_features,n kernel_size=1, stride=1, bias=False))n self.add_module(pool, nn.AvgPool2d(kernel_size=2, stride=2))n

DenseNet

最後實現的DenseNet就是交替連接DenseBlock和Transition（最後一個DenseBlock接池化層和softmax分類器）。

實現代碼如下：

class DenseNet(nn.Module):#121層DenseNetn n def __init__(self, growth_rate=32, block_config=(6, 12, 24, 16),n num_init_features=64, bn_size=4, drop_rate=0, num_classes=1000):nn super(DenseNet, self).__init__()nn # 第一個卷積層n self.features = nn.Sequential(OrderedDict([n (conv0, nn.Conv2d(3, num_init_features, kernel_size=7, stride=2, padding=3, bias=False)),n (norm0, nn.BatchNorm2d(num_init_features)),n (relu0, nn.ReLU(inplace=True)),n (pool0, nn.MaxPool2d(kernel_size=3, stride=2, padding=1)),n ]))nn # 每個denseblockn num_features = num_init_featuresn for i, num_layers in enumerate(block_config):n block = _DenseBlock(num_layers=num_layers, num_input_features=num_features,n bn_size=bn_size, growth_rate=growth_rate, drop_rate=drop_rate)n self.features.add_module(denseblock%d % (i + 1), block)n num_features = num_features + num_layers * growth_raten if i != len(block_config) - 1: #每兩個dense block之間增加一個過渡層n trans = _Transition(num_input_features=num_features, num_output_features=num_features // 2)n self.features.add_module(transition%d % (i + 1), trans)n num_features = num_features // 2nn # batch normn self.features.add_module(norm5, nn.BatchNorm2d(num_features))nn # 分類器n self.classifier = nn.Linear(num_features, num_classes)nn def forward(self, x):n features = self.features(x)n out = F.relu(features, inplace=True)n out = F.avg_pool2d(out, kernel_size=7).view(features.size(0), -1)n out = self.classifier(out)n return outn

這樣就完成了densenet121的pytorch實現代碼。

後記：更多pytorch版本的densenet代碼請參閱pytorch文檔。