Shortcuts

mmcls.models.backbones.hornet 源代码

# Copyright (c) OpenMMLab. All rights reserved.
# Adapted from official impl at https://github.com/raoyongming/HorNet.
try:
    import torch.fft
    fft = True
except ImportError:
    fft = None

import copy
from functools import partial
from typing import Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
from mmcv.cnn.bricks import DropPath

from mmcls.models.backbones.base_backbone import BaseBackbone
from mmcls.registry import MODELS
from ..utils import LayerScale


def get_dwconv(dim, kernel_size, bias=True):
    """build a pepth-wise convolution."""
    return nn.Conv2d(
        dim,
        dim,
        kernel_size=kernel_size,
        padding=(kernel_size - 1) // 2,
        bias=bias,
        groups=dim)


class HorNetLayerNorm(nn.Module):
    """An implementation of LayerNorm of HorNet.

    The differences between HorNetLayerNorm & torch LayerNorm:
        1. Supports two data formats channels_last or channels_first.
    Args:
        normalized_shape (int or list or torch.Size): input shape from an
            expected input of size.
        eps (float): a value added to the denominator for numerical stability.
            Defaults to 1e-5.
        data_format (str): The ordering of the dimensions in the inputs.
            channels_last corresponds to inputs with shape (batch_size, height,
            width, channels) while channels_first corresponds to inputs with
            shape (batch_size, channels, height, width).
            Defaults to 'channels_last'.
    """

    def __init__(self,
                 normalized_shape,
                 eps=1e-6,
                 data_format='channels_last'):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.eps = eps
        self.data_format = data_format
        if self.data_format not in ['channels_last', 'channels_first']:
            raise ValueError(
                'data_format must be channels_last or channels_first')
        self.normalized_shape = (normalized_shape, )

    def forward(self, x):
        if self.data_format == 'channels_last':
            return F.layer_norm(x, self.normalized_shape, self.weight,
                                self.bias, self.eps)
        elif self.data_format == 'channels_first':
            u = x.mean(1, keepdim=True)
            s = (x - u).pow(2).mean(1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.eps)
            x = self.weight[:, None, None] * x + self.bias[:, None, None]
            return x


class GlobalLocalFilter(nn.Module):
    """A GlobalLocalFilter of HorNet.

    Args:
        dim (int): Number of input channels.
        h (int): Height of complex_weight.
            Defaults to 14.
        w (int): Width of complex_weight.
            Defaults to 8.
    """

    def __init__(self, dim, h=14, w=8):
        super().__init__()
        self.dw = nn.Conv2d(
            dim // 2,
            dim // 2,
            kernel_size=3,
            padding=1,
            bias=False,
            groups=dim // 2)
        self.complex_weight = nn.Parameter(
            torch.randn(dim // 2, h, w, 2, dtype=torch.float32) * 0.02)
        self.pre_norm = HorNetLayerNorm(
            dim, eps=1e-6, data_format='channels_first')
        self.post_norm = HorNetLayerNorm(
            dim, eps=1e-6, data_format='channels_first')

    def forward(self, x):
        x = self.pre_norm(x)
        x1, x2 = torch.chunk(x, 2, dim=1)
        x1 = self.dw(x1)

        x2 = x2.to(torch.float32)
        B, C, a, b = x2.shape
        x2 = torch.fft.rfft2(x2, dim=(2, 3), norm='ortho')

        weight = self.complex_weight
        if not weight.shape[1:3] == x2.shape[2:4]:
            weight = F.interpolate(
                weight.permute(3, 0, 1, 2),
                size=x2.shape[2:4],
                mode='bilinear',
                align_corners=True).permute(1, 2, 3, 0)

        weight = torch.view_as_complex(weight.contiguous())

        x2 = x2 * weight
        x2 = torch.fft.irfft2(x2, s=(a, b), dim=(2, 3), norm='ortho')

        x = torch.cat([x1.unsqueeze(2), x2.unsqueeze(2)],
                      dim=2).reshape(B, 2 * C, a, b)
        x = self.post_norm(x)
        return x


class gnConv(nn.Module):
    """A gnConv of HorNet.

    Args:
        dim (int): Number of input channels.
        order (int): Order of gnConv.
            Defaults to 5.
        dw_cfg (dict): The Config for dw conv.
            Defaults to ``dict(type='DW', kernel_size=7)``.
        scale (float): Scaling parameter of gflayer outputs.
            Defaults to 1.0.
    """

    def __init__(self,
                 dim,
                 order=5,
                 dw_cfg=dict(type='DW', kernel_size=7),
                 scale=1.0):
        super().__init__()
        self.order = order
        self.dims = [dim // 2**i for i in range(order)]
        self.dims.reverse()
        self.proj_in = nn.Conv2d(dim, 2 * dim, 1)

        cfg = copy.deepcopy(dw_cfg)
        dw_type = cfg.pop('type')
        assert dw_type in ['DW', 'GF'],\
            'dw_type should be `DW` or `GF`'
        if dw_type == 'DW':
            self.dwconv = get_dwconv(sum(self.dims), **cfg)
        elif dw_type == 'GF':
            self.dwconv = GlobalLocalFilter(sum(self.dims), **cfg)

        self.proj_out = nn.Conv2d(dim, dim, 1)

        self.projs = nn.ModuleList([
            nn.Conv2d(self.dims[i], self.dims[i + 1], 1)
            for i in range(order - 1)
        ])

        self.scale = scale

    def forward(self, x):
        x = self.proj_in(x)
        y, x = torch.split(x, (self.dims[0], sum(self.dims)), dim=1)

        x = self.dwconv(x) * self.scale

        dw_list = torch.split(x, self.dims, dim=1)
        x = y * dw_list[0]

        for i in range(self.order - 1):
            x = self.projs[i](x) * dw_list[i + 1]

        x = self.proj_out(x)

        return x


class HorNetBlock(nn.Module):
    """A block of HorNet.

    Args:
        dim (int): Number of input channels.
        order (int): Order of gnConv.
            Defaults to 5.
        dw_cfg (dict): The Config for dw conv.
            Defaults to ``dict(type='DW', kernel_size=7)``.
        scale (float): Scaling parameter of gflayer outputs.
            Defaults to 1.0.
        drop_path_rate (float): Stochastic depth rate. Defaults to 0.
        use_layer_scale (bool): Whether to use use_layer_scale in HorNet
             block. Defaults to True.
    """

    def __init__(self,
                 dim,
                 order=5,
                 dw_cfg=dict(type='DW', kernel_size=7),
                 scale=1.0,
                 drop_path_rate=0.,
                 use_layer_scale=True):
        super().__init__()
        self.out_channels = dim

        self.norm1 = HorNetLayerNorm(
            dim, eps=1e-6, data_format='channels_first')
        self.gnconv = gnConv(dim, order, dw_cfg, scale)
        self.norm2 = HorNetLayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(dim, 4 * dim)
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(4 * dim, dim)

        if use_layer_scale:
            self.gamma1 = LayerScale(dim, data_format='channels_first')
            self.gamma2 = LayerScale(dim)
        else:
            self.gamma1, self.gamma2 = nn.Identity(), nn.Identity()

        self.drop_path = DropPath(
            drop_path_rate) if drop_path_rate > 0. else nn.Identity()

    def forward(self, x):
        x = x + self.drop_path(self.gamma1(self.gnconv(self.norm1(x))))

        input = x
        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
        x = self.norm2(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        x = self.gamma2(x)
        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)

        x = input + self.drop_path(x)
        return x


[文档]@MODELS.register_module() class HorNet(BaseBackbone): """HorNet backbone. A PyTorch implementation of paper `HorNet: Efficient High-Order Spatial Interactions with Recursive Gated Convolutions <https://arxiv.org/abs/2207.14284>`_ . Inspiration from https://github.com/raoyongming/HorNet Args: arch (str | dict): HorNet architecture. If use string, choose from 'tiny', 'small', 'base' and 'large'. If use dict, it should have below keys: - **base_dim** (int): The base dimensions of embedding. - **depths** (List[int]): The number of blocks in each stage. - **orders** (List[int]): The number of order of gnConv in each stage. - **dw_cfg** (List[dict]): The Config for dw conv. Defaults to 'tiny'. in_channels (int): Number of input image channels. Defaults to 3. drop_path_rate (float): Stochastic depth rate. Defaults to 0. scale (float): Scaling parameter of gflayer outputs. Defaults to 1/3. use_layer_scale (bool): Whether to use use_layer_scale in HorNet block. Defaults to True. out_indices (Sequence[int]): Output from which stages. Default: ``(3, )``. frozen_stages (int): Stages to be frozen (stop grad and set eval mode). -1 means not freezing any parameters. Defaults to -1. with_cp (bool): Use checkpoint or not. Using checkpoint will save some memory while slowing down the training speed. Defaults to False. gap_before_final_norm (bool): Whether to globally average the feature map before the final norm layer. In the official repo, it's only used in classification task. Defaults to True. init_cfg (dict, optional): The Config for initialization. Defaults to None. """ arch_zoo = { **dict.fromkeys(['t', 'tiny'], {'base_dim': 64, 'depths': [2, 3, 18, 2], 'orders': [2, 3, 4, 5], 'dw_cfg': [dict(type='DW', kernel_size=7)] * 4}), **dict.fromkeys(['t-gf', 'tiny-gf'], {'base_dim': 64, 'depths': [2, 3, 18, 2], 'orders': [2, 3, 4, 5], 'dw_cfg': [ dict(type='DW', kernel_size=7), dict(type='DW', kernel_size=7), dict(type='GF', h=14, w=8), dict(type='GF', h=7, w=4)]}), **dict.fromkeys(['s', 'small'], {'base_dim': 96, 'depths': [2, 3, 18, 2], 'orders': [2, 3, 4, 5], 'dw_cfg': [dict(type='DW', kernel_size=7)] * 4}), **dict.fromkeys(['s-gf', 'small-gf'], {'base_dim': 96, 'depths': [2, 3, 18, 2], 'orders': [2, 3, 4, 5], 'dw_cfg': [ dict(type='DW', kernel_size=7), dict(type='DW', kernel_size=7), dict(type='GF', h=14, w=8), dict(type='GF', h=7, w=4)]}), **dict.fromkeys(['b', 'base'], {'base_dim': 128, 'depths': [2, 3, 18, 2], 'orders': [2, 3, 4, 5], 'dw_cfg': [dict(type='DW', kernel_size=7)] * 4}), **dict.fromkeys(['b-gf', 'base-gf'], {'base_dim': 128, 'depths': [2, 3, 18, 2], 'orders': [2, 3, 4, 5], 'dw_cfg': [ dict(type='DW', kernel_size=7), dict(type='DW', kernel_size=7), dict(type='GF', h=14, w=8), dict(type='GF', h=7, w=4)]}), **dict.fromkeys(['b-gf384', 'base-gf384'], {'base_dim': 128, 'depths': [2, 3, 18, 2], 'orders': [2, 3, 4, 5], 'dw_cfg': [ dict(type='DW', kernel_size=7), dict(type='DW', kernel_size=7), dict(type='GF', h=24, w=12), dict(type='GF', h=13, w=7)]}), **dict.fromkeys(['l', 'large'], {'base_dim': 192, 'depths': [2, 3, 18, 2], 'orders': [2, 3, 4, 5], 'dw_cfg': [dict(type='DW', kernel_size=7)] * 4}), **dict.fromkeys(['l-gf', 'large-gf'], {'base_dim': 192, 'depths': [2, 3, 18, 2], 'orders': [2, 3, 4, 5], 'dw_cfg': [ dict(type='DW', kernel_size=7), dict(type='DW', kernel_size=7), dict(type='GF', h=14, w=8), dict(type='GF', h=7, w=4)]}), **dict.fromkeys(['l-gf384', 'large-gf384'], {'base_dim': 192, 'depths': [2, 3, 18, 2], 'orders': [2, 3, 4, 5], 'dw_cfg': [ dict(type='DW', kernel_size=7), dict(type='DW', kernel_size=7), dict(type='GF', h=24, w=12), dict(type='GF', h=13, w=7)]}), } # yapf: disable def __init__(self, arch='tiny', in_channels=3, drop_path_rate=0., scale=1 / 3, use_layer_scale=True, out_indices=(3, ), frozen_stages=-1, with_cp=False, gap_before_final_norm=True, init_cfg=None): super().__init__(init_cfg=init_cfg) if fft is None: raise RuntimeError( 'Failed to import torch.fft. Please install "torch>=1.7".') if isinstance(arch, str): arch = arch.lower() assert arch in set(self.arch_zoo), \ f'Arch {arch} is not in default archs {set(self.arch_zoo)}' self.arch_settings = self.arch_zoo[arch] else: essential_keys = {'base_dim', 'depths', 'orders', 'dw_cfg'} assert isinstance(arch, dict) and set(arch) == essential_keys, \ f'Custom arch needs a dict with keys {essential_keys}' self.arch_settings = arch self.scale = scale self.out_indices = out_indices self.frozen_stages = frozen_stages self.with_cp = with_cp self.gap_before_final_norm = gap_before_final_norm base_dim = self.arch_settings['base_dim'] dims = list(map(lambda x: 2**x * base_dim, range(4))) self.downsample_layers = nn.ModuleList() stem = nn.Sequential( nn.Conv2d(in_channels, dims[0], kernel_size=4, stride=4), HorNetLayerNorm(dims[0], eps=1e-6, data_format='channels_first')) self.downsample_layers.append(stem) for i in range(3): downsample_layer = nn.Sequential( HorNetLayerNorm( dims[i], eps=1e-6, data_format='channels_first'), nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2), ) self.downsample_layers.append(downsample_layer) total_depth = sum(self.arch_settings['depths']) dpr = [ x.item() for x in torch.linspace(0, drop_path_rate, total_depth) ] # stochastic depth decay rule cur_block_idx = 0 self.stages = nn.ModuleList() for i in range(4): stage = nn.Sequential(*[ HorNetBlock( dim=dims[i], order=self.arch_settings['orders'][i], dw_cfg=self.arch_settings['dw_cfg'][i], scale=self.scale, drop_path_rate=dpr[cur_block_idx + j], use_layer_scale=use_layer_scale) for j in range(self.arch_settings['depths'][i]) ]) self.stages.append(stage) cur_block_idx += self.arch_settings['depths'][i] 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.' out_indices = list(out_indices) for i, index in enumerate(out_indices): if index < 0: out_indices[i] = len(self.stages) + index assert 0 <= out_indices[i] <= len(self.stages), \ f'Invalid out_indices {index}.' self.out_indices = out_indices norm_layer = partial( HorNetLayerNorm, eps=1e-6, data_format='channels_first') for i_layer in out_indices: layer = norm_layer(dims[i_layer]) layer_name = f'norm{i_layer}' self.add_module(layer_name, layer) def train(self, mode=True): super(HorNet, self).train(mode) self._freeze_stages() def _freeze_stages(self): for i in range(0, self.frozen_stages + 1): # freeze patch embed m = self.downsample_layers[i] m.eval() for param in m.parameters(): param.requires_grad = False # freeze blocks m = self.stages[i] m.eval() for param in m.parameters(): param.requires_grad = False if i in self.out_indices: # freeze norm m = getattr(self, f'norm{i + 1}') m.eval() for param in m.parameters(): param.requires_grad = False def forward(self, x): outs = [] for i in range(4): x = self.downsample_layers[i](x) if self.with_cp: x = checkpoint.checkpoint_sequential(self.stages[i], len(self.stages[i]), x) else: x = self.stages[i](x) if i in self.out_indices: norm_layer = getattr(self, f'norm{i}') if self.gap_before_final_norm: gap = x.mean([-2, -1], keepdim=True) outs.append(norm_layer(gap).flatten(1)) else: # The output of LayerNorm2d may be discontiguous, which # may cause some problem in the downstream tasks outs.append(norm_layer(x).contiguous()) return tuple(outs)
Read the Docs v: mmcls-1.x
Versions
latest
stable
mmcls-1.x
mmcls-0.x
Downloads
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.