Note
You are reading the documentation for MMClassification 0.x, which will soon be deprecated at the end of 2022. We recommend you upgrade to MMClassification 1.0 to enjoy fruitful new features and better performance brought by OpenMMLab 2.0. Check the installation tutorial, migration tutorial and changelog for more details.
Source code for mmcls.models.backbones.densenet
# Copyright (c) OpenMMLab. All rights reserved.
import math
from itertools import chain
from typing import Sequence
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as cp
from mmcv.cnn.bricks import build_activation_layer, build_norm_layer
from torch.jit.annotations import List
from ..builder import BACKBONES
from .base_backbone import BaseBackbone
class DenseLayer(BaseBackbone):
"""DenseBlock layers."""
def __init__(self,
in_channels,
growth_rate,
bn_size,
norm_cfg=dict(type='BN'),
act_cfg=dict(type='ReLU'),
drop_rate=0.,
memory_efficient=False):
super(DenseLayer, self).__init__()
self.norm1 = build_norm_layer(norm_cfg, in_channels)[1]
self.conv1 = nn.Conv2d(
in_channels,
bn_size * growth_rate,
kernel_size=1,
stride=1,
bias=False)
self.act = build_activation_layer(act_cfg)
self.norm2 = build_norm_layer(norm_cfg, bn_size * growth_rate)[1]
self.conv2 = nn.Conv2d(
bn_size * growth_rate,
growth_rate,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.drop_rate = float(drop_rate)
self.memory_efficient = memory_efficient
def bottleneck_fn(self, xs):
# type: (List[torch.Tensor]) -> torch.Tensor
concated_features = torch.cat(xs, 1)
bottleneck_output = self.conv1(
self.act(self.norm1(concated_features))) # noqa: T484
return bottleneck_output
# todo: rewrite when torchscript supports any
def any_requires_grad(self, x):
# type: (List[torch.Tensor]) -> bool
for tensor in x:
if tensor.requires_grad:
return True
return False
# This decorator indicates to the compiler that a function or method
# should be ignored and replaced with the raising of an exception.
# Here this function is incompatible with torchscript.
@torch.jit.unused # noqa: T484
def call_checkpoint_bottleneck(self, x):
# type: (List[torch.Tensor]) -> torch.Tensor
def closure(*xs):
return self.bottleneck_fn(xs)
# Here use torch.utils.checkpoint to rerun a forward-pass during
# backward in bottleneck to save memories.
return cp.checkpoint(closure, *x)
def forward(self, x): # noqa: F811
# type: (List[torch.Tensor]) -> torch.Tensor
# assert input features is a list of Tensor
assert isinstance(x, list)
if self.memory_efficient and self.any_requires_grad(x):
if torch.jit.is_scripting():
raise Exception('Memory Efficient not supported in JIT')
bottleneck_output = self.call_checkpoint_bottleneck(x)
else:
bottleneck_output = self.bottleneck_fn(x)
new_features = self.conv2(self.act(self.norm2(bottleneck_output)))
if self.drop_rate > 0:
new_features = F.dropout(
new_features, p=self.drop_rate, training=self.training)
return new_features
class DenseBlock(nn.Module):
"""DenseNet Blocks."""
def __init__(self,
num_layers,
in_channels,
bn_size,
growth_rate,
norm_cfg=dict(type='BN'),
act_cfg=dict(type='ReLU'),
drop_rate=0.,
memory_efficient=False):
super(DenseBlock, self).__init__()
self.block = nn.ModuleList([
DenseLayer(
in_channels + i * growth_rate,
growth_rate=growth_rate,
bn_size=bn_size,
norm_cfg=norm_cfg,
act_cfg=act_cfg,
drop_rate=drop_rate,
memory_efficient=memory_efficient) for i in range(num_layers)
])
def forward(self, init_features):
features = [init_features]
for layer in self.block:
new_features = layer(features)
features.append(new_features)
return torch.cat(features, 1)
class DenseTransition(nn.Sequential):
"""DenseNet Transition Layers."""
def __init__(self,
in_channels,
out_channels,
norm_cfg=dict(type='BN'),
act_cfg=dict(type='ReLU')):
super(DenseTransition, self).__init__()
self.add_module('norm', build_norm_layer(norm_cfg, in_channels)[1])
self.add_module('act', build_activation_layer(act_cfg))
self.add_module(
'conv',
nn.Conv2d(
in_channels, out_channels, kernel_size=1, stride=1,
bias=False))
self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2))
[docs]@BACKBONES.register_module()
class DenseNet(BaseBackbone):
"""DenseNet.
A PyTorch implementation of : `Densely Connected Convolutional Networks
<https://arxiv.org/pdf/1608.06993.pdf>`_
Modified from the `official repo
<https://github.com/liuzhuang13/DenseNet>`_
and `pytorch
<https://github.com/pytorch/vision/blob/main/torchvision/models/densenet.py>`_.
Args:
arch (str | dict): The model's architecture. If string, it should be
one of architecture in ``DenseNet.arch_settings``. And if dict, it
should include the following two keys:
- growth_rate (int): Each layer of DenseBlock produce `k` feature
maps. Here refers `k` as the growth rate of the network.
- depths (list[int]): Number of repeated layers in each DenseBlock.
- init_channels (int): The output channels of stem layers.
Defaults to '121'.
in_channels (int): Number of input image channels. Defaults to 3.
bn_size (int): Refers to channel expansion parameter of 1x1
convolution layer. Defaults to 4.
drop_rate (float): Drop rate of Dropout Layer. Defaults to 0.
compression_factor (float): The reduction rate of transition layers.
Defaults to 0.5.
memory_efficient (bool): If True, uses checkpointing. Much more memory
efficient, but slower. Defaults to False.
See `"paper" <https://arxiv.org/pdf/1707.06990.pdf>`_.
norm_cfg (dict): The config dict for norm layers.
Defaults to ``dict(type='BN')``.
act_cfg (dict): The config dict for activation after each convolution.
Defaults to ``dict(type='ReLU')``.
out_indices (Sequence | int): Output from which stages.
Defaults to -1, means the last stage.
frozen_stages (int): Stages to be frozen (all param fixed).
Defaults to 0, which means not freezing any parameters.
init_cfg (dict, optional): Initialization config dict.
"""
arch_settings = {
'121': {
'growth_rate': 32,
'depths': [6, 12, 24, 16],
'init_channels': 64,
},
'169': {
'growth_rate': 32,
'depths': [6, 12, 32, 32],
'init_channels': 64,
},
'201': {
'growth_rate': 32,
'depths': [6, 12, 48, 32],
'init_channels': 64,
},
'161': {
'growth_rate': 48,
'depths': [6, 12, 36, 24],
'init_channels': 96,
},
}
def __init__(self,
arch='121',
in_channels=3,
bn_size=4,
drop_rate=0,
compression_factor=0.5,
memory_efficient=False,
norm_cfg=dict(type='BN'),
act_cfg=dict(type='ReLU'),
out_indices=-1,
frozen_stages=0,
init_cfg=None):
super().__init__(init_cfg=init_cfg)
if isinstance(arch, str):
assert arch in self.arch_settings, \
f'Unavailable arch, please choose from ' \
f'({set(self.arch_settings)}) or pass a dict.'
arch = self.arch_settings[arch]
elif isinstance(arch, dict):
essential_keys = {'growth_rate', 'depths', 'init_channels'}
assert isinstance(arch, dict) and essential_keys <= set(arch), \
f'Custom arch needs a dict with keys {essential_keys}'
self.growth_rate = arch['growth_rate']
self.depths = arch['depths']
self.init_channels = arch['init_channels']
self.act = build_activation_layer(act_cfg)
self.num_stages = len(self.depths)
# check out indices and frozen stages
if isinstance(out_indices, int):
out_indices = [out_indices]
assert isinstance(out_indices, Sequence), \
f'"out_indices" must by a sequence or int, ' \
f'get {type(out_indices)} instead.'
for i, index in enumerate(out_indices):
if index < 0:
out_indices[i] = self.num_stages + index
assert out_indices[i] >= 0, f'Invalid out_indices {index}'
self.out_indices = out_indices
self.frozen_stages = frozen_stages
# Set stem layers
self.stem = nn.Sequential(
nn.Conv2d(
in_channels,
self.init_channels,
kernel_size=7,
stride=2,
padding=3,
bias=False),
build_norm_layer(norm_cfg, self.init_channels)[1], self.act,
nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
# Repetitions of DenseNet Blocks
self.stages = nn.ModuleList()
self.transitions = nn.ModuleList()
channels = self.init_channels
for i in range(self.num_stages):
depth = self.depths[i]
stage = DenseBlock(
num_layers=depth,
in_channels=channels,
bn_size=bn_size,
growth_rate=self.growth_rate,
norm_cfg=norm_cfg,
act_cfg=act_cfg,
drop_rate=drop_rate,
memory_efficient=memory_efficient)
self.stages.append(stage)
channels += depth * self.growth_rate
if i != self.num_stages - 1:
transition = DenseTransition(
in_channels=channels,
out_channels=math.floor(channels * compression_factor),
norm_cfg=norm_cfg,
act_cfg=act_cfg,
)
channels = math.floor(channels * compression_factor)
else:
# Final layers after dense block is just bn with act.
# Unlike the paper, the original repo also put this in
# transition layer, whereas torchvision take this out.
# We reckon this as transition layer here.
transition = nn.Sequential(
build_norm_layer(norm_cfg, channels)[1],
self.act,
)
self.transitions.append(transition)
self._freeze_stages()
[docs] def forward(self, x):
x = self.stem(x)
outs = []
for i in range(self.num_stages):
x = self.stages[i](x)
x = self.transitions[i](x)
if i in self.out_indices:
outs.append(x)
return tuple(outs)
def _freeze_stages(self):
for i in range(self.frozen_stages):
downsample_layer = self.transitions[i]
stage = self.stages[i]
downsample_layer.eval()
stage.eval()
for param in chain(downsample_layer.parameters(),
stage.parameters()):
param.requires_grad = False