1. """dense net in pytorch
    2. """
    3. import torch
    4. import torch.nn as nn
    6. class Bottleneck(nn.Module):
    7.     def __init__(self, in_channels, growth_rate):
    8.         super().__init__()
    9.         inner_channel = 4 * growth_rate
    11.         self.bottle_neck = nn.Sequential(
    12.             nn.BatchNorm2d(in_channels),
    13.             nn.ReLU(inplace=True),
    14.             nn.Conv2d(in_channels, inner_channel, kernel_size=1, bias=False),
    15.             nn.BatchNorm2d(inner_channel),
    16.             nn.ReLU(inplace=True),
    17.             nn.Conv2d(inner_channel, growth_rate, kernel_size=3, padding=1, bias=False)
    18.         )
    20.     def forward(self, x):
    21.         return torch.cat([x, self.bottle_neck(x)], 1)
    23. class Transition(nn.Module):
    24.     def __init__(self, in_channels, out_channels):
    25.         super().__init__()
    27.         self.down_sample = nn.Sequential(
    28.             nn.BatchNorm2d(in_channels),
    29.             nn.Conv2d(in_channels, out_channels, 1, bias=False),
    30.             nn.AvgPool2d(2, stride=2)
    31.         )
    33.     def forward(self, x):
    34.         return self.down_sample(x)
    37. class DenseNet(nn.Module):
    38.     def __init__(self, block, nblocks, growth_rate=12, reduction=0.5, num_class=100):
    39.         super().__init__()
    40.         self.growth_rate = growth_rate
    41.         #扩张率设为32 
    42.         inner_channels = 2 * growth_rate
    43.         #将通道数扩张为64
    44.         self.conv1 = nn.Conv2d(3, inner_channels, kernel_size=3, padding=1, bias=False)
    46.         self.features = nn.Sequential()
    48.         for index in range(len(nblocks) - 1):
    49.             self.features.add_module("dense_block_layer_{}".format(index), self._make_dense_layers(block, inner_channels, nblocks[index]))
    50.             inner_channels += growth_rate * nblocks[index]
    53.             out_channels = int(reduction * inner_channels) # int() will automatic floor the value
    54.             self.features.add_module("transition_layer_{}".format(index), Transition(inner_channels, out_channels))
    55.             inner_channels = out_channels
    57.         self.features.add_module("dense_block{}".format(len(nblocks) - 1), self._make_dense_layers(block, inner_channels, nblocks[len(nblocks)-1]))
    58.         inner_channels += growth_rate * nblocks[len(nblocks) - 1]
    59.         self.features.add_module('bn', nn.BatchNorm2d(inner_channels))
    60.         self.features.add_module('relu', nn.ReLU(inplace=True))
    62.         self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
    64.         self.linear = nn.Linear(inner_channels, num_class)
    66.     def forward(self, x):
    67.         output = self.conv1(x)#64, 32, 32
    68.         output = self.features(output)
    69.         output = self.avgpool(output)
    70.         output = output.view(output.size()[0], -1)
    71.         output = self.linear(output)
    72.         return output
    74.     def _make_dense_layers(self, block, in_channels, nblocks):
    75.         dense_block = nn.Sequential()
    76.         for index in range(nblocks):
    77.             dense_block.add_module('bottle_neck_layer_{}'.format(index), block(in_channels, self.growth_rate))
    78.             in_channels += self.growth_rate
    79.         return dense_block
    81. def densenet121():
    82.     return DenseNet(Bottleneck, [6,12,24,16], growth_rate=32)
    84. def densenet169():
    85.     return DenseNet(Bottleneck, [6,12,32,32], growth_rate=32)
    87. def densenet201():
    88.     return DenseNet(Bottleneck, [6,12,48,32], growth_rate=32)
    90. def densenet161():
    91.     return DenseNet(Bottleneck, [6,12,36,24], growth_rate=48)