#实现DenseNet网络classDenseNet(nn.Module):"DenseNet-BC model"def__init__(self, growth_rate=32, block_config=(6,12,24,16), num_init_features=64,
bn_size=4, compression_rate=0.5, drop_rate=0, num_classes=1000):"""
:param growth_rate: (int) number of filters used in DenseLayer, `k` in the paper
:param block_config: (list of 4 ints) number of layers in each DenseBlock
:param num_init_features: (int) number of filters in the first Conv2d
:param bn_size: (int) the factor using in the bottleneck layer
:param compression_rate: (float) the compression rate used in Transition Layer
:param drop_rate: (float) the drop rate after each DenseLayer
:param num_classes: (int) number of classes for classification
"""super(DenseNet, self).__init__()# first Conv2d
self.features = nn.Sequential(OrderedDict([("conv0", nn.Conv2d(3, num_init_features, kernel_size=7, stride=2, padding=3, bias=False)),("norm0", nn.BatchNorm2d(num_init_features)),("relu0", nn.ReLU(inplace=True)),("pool0", nn.MaxPool2d(3, stride=2, padding=1))]))# DenseBlock
num_features = num_init_features
for i, num_layers inenumerate(block_config):
block = _DenseBlock(num_layers, num_features, bn_size, growth_rate, drop_rate)
self.features.add_module("denseblock%d"%(i +1), block)
num_features += num_layers*growth_rate
if i !=len(block_config)-1:
transition = _Transition(num_features,int(num_features*compression_rate))
self.features.add_module("transition%d"%(i +1), transition)
num_features =int(num_features * compression_rate)# final bn+ReLU
self.features.add_module("norm5", nn.BatchNorm2d(num_features))
self.features.add_module("relu5", nn.ReLU(inplace=True))# classification layer
self.classifier = nn.Linear(num_features, num_classes)# params initializationfor m in self.modules():ifisinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight)elifisinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.bias,0)
nn.init.constant_(m.weight,1)elifisinstance(m, nn.Linear):
nn.init.constant_(m.bias,0)defforward(self, x):
features = self.features(x)
out = F.avg_pool2d(features,7, stride=1).view(features.size(0),-1)
out = self.classifier(out)return out
#选择不同网络参数,就可以实现不同深度的DenseNet,这里实现DenseNet-121网络,而且Pytorch提供了预训练好的网络参数defdensenet121(pretrained=False,**kwargs):"""DenseNet121"""
model = DenseNet(num_init_features=64, growth_rate=32, block_config=(6,12,24,16),**kwargs)if pretrained:# '.'s are no longer allowed in module names, but pervious _DenseLayer# has keys 'norm.1', 'relu.1', 'conv.1', 'norm.2', 'relu.2', 'conv.2'.# They are also in the checkpoints in model_urls. This pattern is used# to find such keys.
pattern = re.compile(
r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')
state_dict = model_zoo.load_url(model_urls['densenet121'])for key inlist(state_dict.keys()):
res = pattern.match(key)if res:
new_key = res.group(1)+ res.group(2)
state_dict[new_key]= state_dict[key]del state_dict[key]
model.load_state_dict(state_dict)return model