Shortcuts

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
[docs] def train(self, mode=True): super(DenseNet, self).train(mode) self._freeze_stages()
Read the Docs v: latest
Versions
master
latest
1.x
dev-1.x
Downloads
html
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.