import sys
_module = sys.modules[__name__]
del sys
setup = _module
test_dataloader = _module
test_imagenet = _module
albumentation_vs_torchvision = _module
build = _module
augment = _module
color_jitter = _module
crop = _module
dropout = _module
flip = _module
pad = _module
resize = _module
rotate = _module
ghostnet = _module
mobilenet = _module
resnet = _module
shufflenet = _module
test_basicblock = _module
test_bottleneck = _module
test_ghost_backbone = _module
test_mnasnet_beckbone = _module
test_mnasnet_unit = _module
test_mobilenetv1_backbone = _module
test_mobilenetv1_block = _module
test_mobilenetv2_backbone = _module
test_mobilenetv2_block = _module
test_mobilenetv2_inverted_residual = _module
test_mobilenetv3_backbone = _module
test_mobilenetv3_unit = _module
test_repvgg_backbone = _module
test_resnest_block = _module
test_resnet3d_backbone = _module
test_resnet3d_basicblock = _module
test_resnet3d_bottleneck = _module
test_resnet_backbone = _module
test_resnet_d_backbone = _module
test_shufflenetv1_backbone = _module
test_shufflenetv1_unit = _module
test_shufflenetv2_backbone = _module
test_shufflenetv2_unit = _module
test_sknet_block = _module
label_smoothing_loss = _module
test_general_head_2d = _module
test_general_head_3d = _module
test_mobilenet_v3_head = _module
test_asymmetric_convolution_block = _module
test_batchnorm2d = _module
test_dbblock = _module
test_global_context_block = _module
test_groupnorm = _module
test_layernorm = _module
test_non_local_embedded_gaussian = _module
test_place_holder = _module
test_repvgg_block = _module
test_selective_kernel_conv2d = _module
test_simplified_non_local_embedded_gaussian = _module
test_split_attention_conv2d = _module
test_squeeze_and_excitation_block = _module
test_ghostnet = _module
test_mnasnet = _module
test_mobilenet_v1 = _module
test_mobilenet_v2 = _module
test_mobilenet_v3 = _module
test_regvgg = _module
test_resnest = _module
test_resnet = _module
test_shufflenetv1 = _module
test_shufflenetv2 = _module
test_sknet = _module
test_group_weight = _module
test_prefetcher = _module
benchmarks = _module
compute_mean_std = _module
moskomule_se_resnet_to_zcls_se_resnet = _module
official_ghostnet_to_zcls_ghostnet = _module
official_repvgg_to_zcls_repvgg = _module
official_resnest_to_zcls_resnest = _module
official_shufflenet_v1_to_zcls_shufflenet_v1 = _module
syt2_sknet_to_zcls_sknet = _module
torchvision_mobilenet_to_zcls_mobilenet = _module
torchvision_resnet_to_zcls_resnet = _module
torchvision_shufflenet_to_zcls_shufflenet = _module
model_fuse = _module
add_md5_for_pths = _module
test = _module
train = _module
zoom = _module
zcls = _module
assets = _module
config = _module
configs = _module
custom_config = _module
dataloader = _module
dataset = _module
defaults = _module
lr_scheduler = _module
model = _module
optimizer = _module
transform = _module
key_word = _module
data = _module
build = _module
build = _module
datasets = _module
build = _module
cifar = _module
evaluator = _module
base_evaluator = _module
general_evaluator = _module
fashionmnist = _module
general_dataset = _module
general_dataset_v2 = _module
imagenet = _module
lmdb_dataset = _module
lmdb_imagenet = _module
mp_dataset = _module
util = _module
samplers = _module
distributed_sampler = _module
transforms = _module
realization = _module
autoaugment = _module
center_crop = _module
coarse_dropout = _module
horizontal_flip = _module
normalize = _module
random_crop = _module
square_pad = _module
to_tensor = _module
vertical_flip = _module
engine = _module
inference = _module
trainer = _module
act_helper = _module
attention_helper = _module
backbones = _module
ghost_backbone = _module
ghost_bottleneck = _module
ghost_module = _module
misc = _module
mnasnet_backbone = _module
mnasnet_unit = _module
mobilenetv1_backbone = _module
mobilenetv1_block = _module
mobilenetv2_backbone = _module
mobilenetv2_block = _module
mobilenetv2_inverted_residual = _module
mobilenetv3_backbone = _module
mobilenetv3_unit = _module
basicblock = _module
bottleneck = _module
resnest_block = _module
resnet3d_backbone = _module
resnet3d_basicblock = _module
resnet3d_bottleneck = _module
resnet_backbone = _module
resnet_d_backbone = _module
sknet_block = _module
shufflenetv1_backbone = _module
shufflenetv1_unit = _module
shufflenetv2_backbone = _module
shufflenetv2_unit = _module
vgg = _module
repvgg_backbone = _module
conv_helper = _module
criterions = _module
crossentropy_loss = _module
label_smoothing_loss = _module
heads = _module
general_head_2d = _module
general_head_3d = _module
mobilenetv3_head = _module
init_helper = _module
layers = _module
acb_util = _module
asymmetric_convolution_block = _module
batch_norm_2d = _module
dbb_transforms = _module
dbb_util = _module
diverse_branch_block = _module
global_context_block = _module
group_norm = _module
group_norm_wrapper = _module
layer_norm = _module
non_local_embedded_gaussian = _module
place_holder = _module
repvgg_block = _module
repvgg_util = _module
selective_kernel_conv2d = _module
simplified_non_local_embedded_gaussian = _module
split_attention_conv2d = _module
squeeze_and_excitation_block = _module
misc = _module
norm_helper = _module
recognizers = _module
base_recognizer = _module
build = _module
mnasnet = _module
mobilenetv1 = _module
mobilenetv2 = _module
mobilenetv3 = _module
torchvision_mnasnet = _module
torchvision_mobilenetv2 = _module
torchvision_mobilenetv3 = _module
official_resnest = _module
resnet3d = _module
torchvision_resnet = _module
shufflenetv1 = _module
shufflenetv2 = _module
torchvision_sfv2 = _module
repvgg = _module
registry = _module
optim = _module
lr_schedulers = _module
build = _module
cosine_annealing_lr = _module
gradual_warmup = _module
multistep_lr = _module
optimizers = _module
adam = _module
build = _module
rmsprop = _module
sgd = _module
checkpoint = _module
collect_env = _module
cutmix = _module
distributed = _module
logging = _module
metric_logger = _module
metrics = _module
misc = _module
mixup = _module
multiprocessing = _module
parser = _module
precise_bn = _module
prefetcher = _module

from _paritybench_helpers import _mock_config, patch_functional
from unittest.mock import mock_open, MagicMock
from torch.autograd import Function
from torch.nn import Module
import abc, collections, copy, enum, functools, inspect, itertools, logging, math, matplotlib, numbers, numpy, pandas, queue, random, re, scipy, sklearn, string, tensorflow, time, torch, torchaudio, torchtext, torchvision, types, typing, uuid, warnings
import numpy as np
from torch import Tensor
patch_functional()
open = mock_open()
yaml = logging = sys = argparse = MagicMock()
ArgumentParser = argparse.ArgumentParser
_global_config = args = argv = cfg = config = params = _mock_config()
argparse.ArgumentParser.return_value.parse_args.return_value = _global_config
yaml.load.return_value = _global_config
sys.argv = _global_config
__version__ = '1.0.0'
xrange = range
wraps = functools.wraps


import torch


import torch.nn as nn


import numpy as np


import torch.nn.functional as F


from torch.nn.modules.loss import _WeightedLoss


from torchvision.models import resnet50


import time


from torchvision.datasets.cifar import CIFAR100


from torch.utils.data import DataLoader


import torchvision.transforms as transforms


from torch.utils.data.distributed import DistributedSampler


from torch.utils.data import RandomSampler


from torch.utils.data import SequentialSampler


from torch.utils.data import Dataset


from torchvision.datasets import CIFAR10


from torchvision.datasets import CIFAR100


from abc import ABCMeta


from abc import abstractmethod


import torchvision.datasets as datasets


import logging


from torch.utils.data import IterableDataset


import math


from typing import TypeVar


from typing import Optional


from typing import Iterator


from torch.utils.data import Sampler


import torch.distributed as dist


from torch.nn.parallel import DistributedDataParallel


from torch.cuda.amp import GradScaler


from torch.cuda.amp import autocast


import copy


from abc import ABC


from torchvision.models.resnet import model_urls


from torch.nn import init


from torch import Tensor


from torch.utils import model_zoo


from functools import partial


from torch.nn.modules.module import T


from torch.nn.parallel import DistributedDataParallel as DDP


from torchvision.models.mnasnet import mnasnet0_5


from torchvision.models.mnasnet import mnasnet0_75


from torchvision.models.mnasnet import mnasnet1_0


from torchvision.models.mnasnet import mnasnet1_3


from torchvision.models import mobilenet_v2


from torchvision.models import mobilenet_v3_large


from torchvision.models import mobilenet_v3_small


from torchvision.models.resnet import resnet18


from torchvision.models.resnet import resnet34


from torchvision.models.resnet import resnet50


from torchvision.models.resnet import resnet101


from torchvision.models.resnet import resnet152


from torchvision.models.resnet import resnext50_32x4d


from torchvision.models.resnet import resnext101_32x8d


from torchvision.models.shufflenetv2 import shufflenet_v2_x0_5


from torchvision.models.shufflenetv2 import shufflenet_v2_x1_0


from torchvision.models.shufflenetv2 import shufflenet_v2_x1_5


from torchvision.models.shufflenetv2 import shufflenet_v2_x2_0


from torch.optim.optimizer import Optimizer


import torch.optim as optim


from torch.optim.lr_scheduler import _LRScheduler


from torch.optim.lr_scheduler import ReduceLROnPlateau


import re


from collections import defaultdict


import torchvision


import functools


from collections import deque


import torch.multiprocessing as mp


import itertools


from typing import Any


from typing import Iterable


from typing import List


from typing import Tuple


from typing import Type


from torch import nn


class TestA(nn.Module):

    def __init__(self):
        super(TestA, self).__init__()
        self.conv1 = nn.Conv2d(3, 5, kernel_size=1, stride=1)
        self.bn1 = nn.BatchNorm2d(5)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Sequential(nn.Conv2d(5, 10, kernel_size=1, stride=1), nn.BatchNorm2d(10), nn.ReLU(inplace=True))
        self.fc = nn.Linear(10, 20)


class GhostModule(nn.Module):

    def __init__(self, in_channels, out_channels, ratio=2, primary_kernel_size=1, cheap_kernel_size=3, stride=1, is_act=True) ->None:
        """
        Block = primary_conv + cheap_operation
        :param in_channels: 输入通道数
        :param out_channels: 输出通道数
        :param ratio: 创建内在特征图比率
        :param primary_kernel_size: primary卷积核大小
        :param cheap_kernel_size: cheap_operation深度卷积核大小
        :param stride: primary步长
        :param is_act: 是否执行激活函数
        """
        super(GhostModule, self).__init__()
        self.out_channels = out_channels
        init_channels = math.ceil(out_channels / ratio)
        new_channels = init_channels * (ratio - 1)
        primary_padding = primary_kernel_size // 2
        self.primary_conv = nn.Sequential(nn.Conv2d(in_channels, init_channels, kernel_size=(primary_kernel_size, primary_kernel_size), stride=(stride, stride), padding=(primary_padding, primary_padding), bias=False), nn.BatchNorm2d(init_channels), nn.ReLU(inplace=True) if is_act else nn.Sequential())
        cheap_padding = cheap_kernel_size // 2
        self.cheap_operation = nn.Sequential(nn.Conv2d(init_channels, new_channels, kernel_size=(cheap_kernel_size, cheap_kernel_size), stride=(1, 1), padding=(cheap_padding, cheap_padding), groups=init_channels, bias=False), nn.BatchNorm2d(new_channels), nn.ReLU(inplace=True) if is_act else nn.Sequential())

    def forward(self, x):
        x1 = self.primary_conv(x)
        x2 = self.cheap_operation(x1)
        out = torch.cat([x1, x2], dim=1)
        return out[:, :self.out_channels, :, :]


class _GlobalContextBlockND(nn.Module):
    """
    refer to [context_block.py](https://github.com/xvjiarui/GCNet/blob/master/mmdet/ops/gcb/context_block.py)
    """

    def __init__(self, in_channels, reduction=16, dimension=2):
        """
        :param in_channels:
        :param reduction:
        :param dimension:
        """
        super(_GlobalContextBlockND, self).__init__()
        assert dimension in [1, 2, 3]
        self.in_channels = in_channels
        self.reduction = reduction
        self.dimension = dimension
        if self.dimension == 1:
            self.conv_layer = nn.Conv1d
        elif self.dimension == 2:
            self.conv_layer = nn.Conv2d
        else:
            self.conv_layer = nn.Conv3d
        self._construct_gc()
        self.init_weights()

    def _construct_gc(self):
        self.w_k = self.conv_layer(self.in_channels, 1, kernel_size=1, stride=1, padding=0)
        reduction_channel = self.in_channels // self.reduction
        self.w_v1 = self.conv_layer(self.in_channels, reduction_channel, kernel_size=1, stride=1, padding=0)
        self.ln = nn.LayerNorm([reduction_channel, *([1] * self.dimension)])
        self.relu = nn.ReLU(inplace=True)
        self.w_v2 = self.conv_layer(reduction_channel, self.in_channels, kernel_size=1, stride=1, padding=0)

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')

    def forward(self, x):
        """
        :param x: (N, C, **)
        """
        self._check_input_dim(x)
        identity = x
        N, C = x.shape[:2]
        input_x = x.view(N, C, -1).unsqueeze(1)
        context_mask = self.w_k(x).view(N, 1, -1)
        context_mask = nn.functional.softmax(context_mask, dim=2).unsqueeze(-1)
        context = torch.matmul(input_x, context_mask).reshape(N, C).reshape(N, C, *([1] * self.dimension))
        out = self.w_v1(context)
        out = self.ln(out)
        out = self.relu(out)
        transform = self.w_v2(out)
        z = transform + identity
        return z

    def _check_input_dim(self, input):
        raise NotImplementedError


class GlobalContextBlock2D(_GlobalContextBlockND):

    def __init__(self, in_channels, reduction=16, dimension=2):
        super().__init__(in_channels, reduction, dimension)

    def _check_input_dim(self, input):
        if input.dim() != 4:
            raise ValueError('expected 4D input (got {}D input)'.format(input.dim()))


class _NonLocalNDEmbeddedGaussian(nn.Module):
    """
    refer to
    1. https://github.com/AlexHex7/Non-local_pytorch/blob/master/lib/non_local_embedded_gaussian.py
    2. https://github.com/facebookresearch/SlowFast/blob/master/slowfast/models/nonlocal_helper.py
    """

    def __init__(self, in_channels, inner_channels=None, dimension=2, with_pool=True, norm_layer=None, zero_init_final_norm=True):
        super(_NonLocalNDEmbeddedGaussian, self).__init__()
        assert dimension in [1, 2, 3]
        self.in_channels = in_channels
        self.inner_channels = inner_channels
        self.dimension = dimension
        self.with_pool = with_pool
        self.norm_layer = norm_layer
        self.zero_init_final_norm = zero_init_final_norm
        if self.inner_channels is None:
            self.inner_channels = self.in_channels // 2
            if self.inner_channels == 0:
                self.inner_channels = 1
        if self.dimension == 1:
            self.conv_layer = nn.Conv1d
            if self.norm_layer is None:
                self.norm_layer = nn.BatchNorm1d
            if with_pool:
                self.pool = nn.MaxPool1d(kernel_size=2)
        elif self.dimension == 2:
            self.conv_layer = nn.Conv2d
            if self.norm_layer is None:
                self.norm_layer = nn.BatchNorm2d
            if with_pool:
                self.pool = nn.MaxPool2d(kernel_size=(2, 2))
        else:
            self.conv_layer = nn.Conv3d
            if self.norm_layer is None:
                self.norm_layer = nn.BatchNorm3d
            if with_pool:
                self.pool = nn.MaxPool3d(kernel_size=(1, 2, 2))
        self._construct_nonlocal()
        self.init_weights()

    def _construct_nonlocal(self):
        self.w_theta = self.conv_layer(self.in_channels, self.inner_channels, kernel_size=1, stride=1, padding=0)
        self.w_phi = self.conv_layer(self.in_channels, self.inner_channels, kernel_size=1, stride=1, padding=0)
        self.w_g = self.conv_layer(self.in_channels, self.inner_channels, kernel_size=1, stride=1, padding=0)
        if self.with_pool:
            self.w_phi = nn.Sequential(self.w_phi, self.pool)
            self.w_g = nn.Sequential(self.w_g, self.pool)
        self.w_z = nn.Sequential(self.conv_layer(self.inner_channels, self.in_channels, kernel_size=1, stride=1, padding=0), self.norm_layer(self.in_channels))

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.GroupNorm)):
                if self.zero_init_final_norm:
                    nn.init.constant_(m.weight, 0)
                else:
                    nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        """
        :param x: (N, C, **)
        """
        self._check_input_dim(x)
        identity = x
        N, C = x.shape[:2]
        theta_x = self.w_theta(x).view(N, self.inner_channels, -1)
        phi_x = self.w_phi(x).view(N, self.inner_channels, -1)
        f = torch.einsum('nct,ncp->ntp', (theta_x, phi_x))
        f_div_C = f * self.inner_channels ** -0.5
        f_div_C = nn.functional.softmax(f_div_C, dim=2)
        g_x = self.w_g(x).view(N, self.inner_channels, -1)
        y = torch.einsum('ntg,ncg->nct', (f_div_C, g_x))
        y = y.view(N, self.inner_channels, *x.shape[2:])
        w_y = self.w_z(y).contiguous()
        z = w_y + identity
        return z

    def _check_input_dim(self, input):
        raise NotImplementedError


class NonLocal2DEmbeddedGaussian(_NonLocalNDEmbeddedGaussian):

    def __init__(self, in_channels, inner_channels=None, dimension=2, with_pool=True, norm_layer=None, zero_init_final_norm=True):
        super().__init__(in_channels, inner_channels, dimension, with_pool, norm_layer, zero_init_final_norm)

    def _check_input_dim(self, input):
        if input.dim() != 4:
            raise ValueError('expected 4D input (got {}D input)'.format(input.dim()))


class _SimplifiedNonLocalNDEmbeddedGaussian(nn.Module):

    def __init__(self, in_channels, dimension=2, norm_layer=None, zero_init_final_norm=True):
        """
        refer to [context_block.py](https://github.com/xvjiarui/GCNet/blob/master/mmdet/ops/gcb/context_block.py)
        :param in_channels: Input channels
        :param dimension: Data dimension
        :param norm_layer: Normalized layer type
        :param zero_init_final_norm: Initialize the last BN layer at zero to ensure the consistent connection of the initial model
        """
        super(_SimplifiedNonLocalNDEmbeddedGaussian, self).__init__()
        assert dimension in [1, 2, 3]
        self.in_channels = in_channels
        self.dimension = dimension
        self.norm_layer = norm_layer
        self.zero_init_final_norm = zero_init_final_norm
        if self.dimension == 1:
            self.conv_layer = nn.Conv1d
            if self.norm_layer is None:
                self.norm_layer = nn.BatchNorm1d
        elif self.dimension == 2:
            self.conv_layer = nn.Conv2d
            if self.norm_layer is None:
                self.norm_layer = nn.BatchNorm2d
        else:
            self.conv_layer = nn.Conv3d
            if self.norm_layer is None:
                self.norm_layer = nn.BatchNorm3d
        self._construct_nonlocal()
        self.init_weights()

    def _construct_nonlocal(self):
        self.w_k = self.conv_layer(self.in_channels, 1, kernel_size=1, stride=1, padding=0)
        self.w_v = nn.Sequential(self.conv_layer(self.in_channels, self.in_channels, kernel_size=1, stride=1, padding=0), self.norm_layer(self.in_channels))

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.GroupNorm)):
                if self.zero_init_final_norm:
                    nn.init.constant_(m.weight, 0)
                else:
                    nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        """
        :param x: (N, C, **)
        """
        self._check_input_dim(x)
        identity = x
        N, C = x.shape[:2]
        input_x = x.view(N, C, -1).unsqueeze(1)
        context_mask = self.w_k(x).view(N, 1, -1).unsqueeze(-1)
        context_mask = nn.functional.softmax(context_mask, dim=2)
        context = torch.matmul(input_x, context_mask).reshape(N, C).reshape(N, C, *([1] * self.dimension))
        transform = self.w_v(context)
        z = transform + identity
        return z

    def _check_input_dim(self, input):
        raise NotImplementedError


class SimplifiedNonLocal2DEmbeddedGaussian(_SimplifiedNonLocalNDEmbeddedGaussian):

    def __init__(self, in_channels, dimension=2, norm_layer=None, zero_init_final_norm=True):
        super().__init__(in_channels, dimension, norm_layer, zero_init_final_norm)

    def _check_input_dim(self, input):
        if input.dim() != 4:
            raise ValueError('expected 4D input (got {}D input)'.format(input.dim()))


def get_sigmoid(sigmoid_type):
    if sigmoid_type == 'Sigmoid':
        return nn.Sigmoid
    elif sigmoid_type == 'HSigmoid':
        return nn.Hardsigmoid
    else:
        raise IOError(f"{sigmoid_type} doesn't exists")


def round_to_multiple_of(val, divisor, round_up_bias=0.9):
    """ Asymmetric rounding to make `val` divisible by `divisor`. With default
    bias, will round up, unless the number is no more than 10% greater than the
    smaller divisible value, i.e. (83, 8) -> 80, but (84, 8) -> 88. """
    assert 0.0 < round_up_bias < 1.0
    new_val = max(divisor, int(val + divisor / 2) // divisor * divisor)
    return new_val if new_val >= round_up_bias * val else new_val + divisor


class _SqueezeAndExcitationBlockND(nn.Module):

    def __init__(self, in_channels, reduction=16, dimension=2, sigmoid_type='Sigmoid', bias=False, is_round=False, round_nearest=8):
        """
        Squeeze-and-Excitation Block
        refer to
        [se_module.py](https://github.com/moskomule/senet.pytorch/blob/master/senet/se_module.py)
        [如何评价Momenta ImageNet 2017夺冠架构SENet?](https://www.zhihu.com/question/63460684)
        :param in_channels:
        :param reduction:
        :param dimension:
        :param sigmoid_type:
        """
        super(_SqueezeAndExcitationBlockND, self).__init__()
        assert dimension in [1, 2, 3]
        assert in_channels % reduction == 0, f'in_channels = {in_channels}, reduction = {reduction}'
        inner_channel = in_channels // reduction
        if is_round:
            inner_channel = round_to_multiple_of(inner_channel, round_nearest)
        if dimension == 1:
            self.squeeze = nn.AdaptiveAvgPool1d(1)
        elif dimension == 2:
            self.squeeze = nn.AdaptiveAvgPool2d((1, 1))
        else:
            self.squeeze = nn.AdaptiveAvgPool3d((1, 1, 1))
        sigmoid_layer = get_sigmoid(sigmoid_type)
        self.excitation = nn.Sequential(nn.Linear(in_channels, inner_channel, bias=bias), nn.ReLU(inplace=True), nn.Linear(inner_channel, in_channels, bias=bias), sigmoid_layer())
        self.init_weights()

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)

    def forward(self, x):
        """
        :param x: (N, C, **)
        """
        self._check_input_dim(x)
        N, C = x.shape[:2]
        out = self.squeeze(x)
        out_shape = out.shape
        out = self.excitation(out.view(N, C)).view(out_shape)
        scale = x * out.expand_as(x)
        return scale

    def _check_input_dim(self, input):
        raise NotImplementedError


class SqueezeAndExcitationBlock2D(_SqueezeAndExcitationBlockND):

    def __init__(self, in_channels, reduction=16, dimension=2, **kwargs):
        super().__init__(in_channels, reduction, dimension, **kwargs)

    def _check_input_dim(self, input):
        if input.dim() != 4:
            raise ValueError('expected 4D input (got {}D input)'.format(input.dim()))


def make_attention_block(in_planes, reduction, attention_type, **kwargs):
    if attention_type == 'GlobalContextBlock2D':
        return GlobalContextBlock2D(in_channels=in_planes, reduction=reduction)
    elif attention_type == 'SqueezeAndExcitationBlock2D':
        return SqueezeAndExcitationBlock2D(in_channels=in_planes, reduction=reduction, **kwargs)
    elif attention_type == 'NonLocal2DEmbeddedGaussian':
        return NonLocal2DEmbeddedGaussian(in_channels=in_planes)
    elif attention_type == 'SimplifiedNonLocal2DEmbeddedGaussian':
        return SimplifiedNonLocal2DEmbeddedGaussian(in_channels=in_planes)
    else:
        raise ValueError('no matching type')


class GhostBottleneck(nn.Module):

    def __init__(self, in_channels, mid_channels, out_channels, stride=1, kernel_size=3, with_attention=False, reduction=4, attention_type='SqueezeAndExcitationBlock2D', sigmoid_type='HSigmoid'):
        """
        when stride=1, bottleneck=Add(
                                    Input,
                                    BN(GhostModule(
                                        ReLU(BN(
                                            GhostModule(Input)
                                        ))
                                    ))
                                    )
        when stride=2, bottleneck=Add(
                                    Input,
                                    BN(GhostModule(
                                        BN(DWConv(
                                            ReLU(BN(
                                                GhostModule(Input)
                                            ))
                                        ))
                                    ))
                                    )
        if assert squeeze-and-excitation module, use it before ghost2
        :param in_channels: 输入通道数
        :param mid_channels: 膨胀通道数
        :param out_channels: 输出通道数
        :param stride: 步长，是否执行下采样操作
        :param kernel_size: 深度卷积核大小
        :param with_attention: 是否执行注意力模块
        :param reduction: 衰减率
        :param attention_type: 注意力类型
        """
        super(GhostBottleneck, self).__init__()
        self.ghost1 = GhostModule(in_channels, mid_channels, is_act=True)
        self.stride = stride
        if self.stride > 1:
            padding = (kernel_size - 1) // 2
            self.conv_dw = nn.Conv2d(mid_channels, mid_channels, (kernel_size, kernel_size), stride=(stride, stride), padding=(padding, padding), groups=mid_channels, bias=False)
            self.bn_dw = nn.BatchNorm2d(mid_channels)
        if with_attention:
            self.se = make_attention_block(mid_channels, reduction, attention_type, sigmoid_type=sigmoid_type, is_round=True, round_nearest=4, bias=True)
        self.with_attention = with_attention
        self.ghost2 = GhostModule(mid_channels, out_channels, is_act=False)
        if in_channels == out_channels and self.stride == 1:
            self.shortcut = nn.Sequential()
        else:
            padding = (kernel_size - 1) // 2
            self.shortcut = nn.Sequential(nn.Conv2d(in_channels, in_channels, (kernel_size, kernel_size), stride=(stride, stride), padding=(padding, padding), groups=in_channels, bias=False), nn.BatchNorm2d(in_channels), nn.Conv2d(in_channels, out_channels, (1, 1), stride=(1, 1), padding=(0, 0), bias=False), nn.BatchNorm2d(out_channels))

    def forward(self, x):
        residual = x
        x = self.ghost1(x)
        if self.stride > 1:
            x = self.conv_dw(x)
            x = self.bn_dw(x)
        if self.with_attention:
            x = self.se(x)
        x = self.ghost2(x)
        x += self.shortcut(residual)
        return x


arch_setting = [[[3, 16, 16, 0, 1]], [[3, 48, 24, 0, 2]], [[3, 72, 24, 0, 1]], [[5, 72, 40, 4, 2]], [[5, 120, 40, 4, 1]], [[3, 240, 80, 0, 2]], [[3, 200, 80, 0, 1], [3, 184, 80, 0, 1], [3, 184, 80, 0, 1], [3, 480, 112, 4, 1], [3, 672, 112, 4, 1]], [[5, 672, 160, 4, 2]], [[5, 960, 160, 0, 1], [5, 960, 160, 4, 1], [5, 960, 160, 0, 1], [5, 960, 160, 4, 1]]]


def make_stem(in_channels, base_channels, groups, with_attention, reduction, attention_type, attention_bias, conv_layer, act_layer):
    """

    :param in_channels: 输入通道数
    :param base_channels: 卷积层输出通道数
    :param groups: 分组数
    :param conv_layer: 卷积层
    :param act_layer: 激活层类型
    :return:
    """
    return nn.Sequential(conv_layer(in_channels=in_channels, out_channels=base_channels, kernel_size=3, stride=2, padding=1, groups=groups, bias=True), make_attention_block(base_channels, reduction, attention_type, bias=attention_bias) if with_attention else nn.Identity(), act_layer(inplace=True))


class GhostBackbone(nn.Module):

    def __init__(self, in_channels=3, base_channels=16, width_multiplier=1.0, round_nearest=4, attention_type='SqueezeAndExcitationBlock2D', sigmoid_type='HSigmoid', arch_configs=None, block_layer=None, conv_layer=None, norm_layer=None, act_layer=None, **kwargs):
        """
        :param in_channels: 输入通道数
        :param base_channels: 基础通道数
        :param width_multiplier: 宽度乘法器
        :param round_nearest: 设置每一层通道数均为4的倍数
        :param attention_type: 注意力模块类型
        :param sigmoid_type: 作用于注意力模块
        :param block_layer: 块类型
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        :param zero_init_residual: 零初始化残差连接
        :param kwargs: 其他参数
        """
        super(GhostBackbone, self).__init__()
        if arch_configs is None:
            arch_configs = copy.deepcopy(arch_setting)
        if block_layer is None:
            block_layer = GhostBottleneck
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        base_channels = round_to_multiple_of(base_channels * width_multiplier, round_nearest)
        last_in_channels = round_to_multiple_of(arch_configs[-1][-1][2] * width_multiplier, round_nearest)
        last_out_channels = round_to_multiple_of(arch_configs[-1][-1][1] * width_multiplier, round_nearest)
        self.first_stem, self.last_stem = make_stem(in_channels, base_channels, last_in_channels, last_out_channels, conv_layer, norm_layer, act_layer)
        input_channels = base_channels
        stages = []
        for cfg in arch_configs:
            layers = []
            for kernel_size, expansion_size, out_channels, attention_reduction_rate, stride in cfg:
                hidden_channels = round_to_multiple_of(expansion_size * width_multiplier, round_nearest)
                output_channels = round_to_multiple_of(out_channels * width_multiplier, round_nearest)
                layers.append(block_layer(input_channels, hidden_channels, output_channels, stride=stride, kernel_size=kernel_size, with_attention=attention_reduction_rate > 1, reduction=attention_reduction_rate, attention_type=attention_type, sigmoid_type=sigmoid_type))
                input_channels = output_channels
            stages.append(nn.Sequential(*layers))
        self.blocks = nn.Sequential(*stages)

    def forward(self, x):
        x = self.first_stem(x)
        x = self.blocks(x)
        x = self.last_stem(x)
        return x


BN_MOMENTUM = 1 - 0.9997


class MNASNetUint(nn.Module, ABC):

    def __init__(self, in_planes, out_planes, stride=1, kernel_size=3, expansion_rate=1, with_attention=True, reduction=4, attention_type='SqueezeAndExcitationBlock2D', conv_layer=None, norm_layer=None, act_layer=None):
        """
        整合了MBConv3(3x3/5x5/SE)、MBConv6(3x3/5x5/SE)、SepConv(k3x3)实现
        :param in_planes: 输入通道数
        :param out_planes: 输出通道数
        :param stride: 步长
        :param kernel_size: 卷积核大小
        :param expansion_rate: 膨胀因子，适用于mbv2的反向残差块
        :param with_attention: 是否使用注意力模块
        :param reduction: 衰减率
        :param attention_type: 注意力模块类型
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        """
        super(MNASNetUint, self).__init__()
        assert isinstance(expansion_rate, int) and expansion_rate >= 1
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        self.with_expansion = expansion_rate > 1
        self.with_attention = with_attention
        inner_planes = in_planes * expansion_rate
        if self.with_expansion:
            self.expansion = nn.Sequential(conv_layer(in_planes, inner_planes, kernel_size=1, stride=1, padding=0, bias=False), norm_layer(inner_planes, momentum=BN_MOMENTUM), act_layer(inplace=True))
        if kernel_size == 3:
            padding = 1
        elif kernel_size == 5:
            padding = 2
        else:
            padding = 0
        self.conv1 = conv_layer(inner_planes, inner_planes, kernel_size=kernel_size, stride=stride, padding=padding, bias=False, groups=inner_planes)
        self.norm1 = norm_layer(inner_planes, momentum=BN_MOMENTUM)
        self.act = act_layer(inplace=True)
        self.conv2 = conv_layer(inner_planes, out_planes, kernel_size=1, stride=1, padding=0, bias=False)
        self.norm2 = norm_layer(out_planes, momentum=BN_MOMENTUM)
        if self.with_attention:
            self.attention = make_attention_block(inner_planes, reduction, attention_type)
        self.apply_residual = in_planes == out_planes and stride == 1

    def forward(self, x):
        identity = x
        if self.with_expansion:
            x = self.expansion(x)
        out = self.conv1(x)
        out = self.norm1(out)
        out = self.act(out)
        if self.with_attention:
            out = self.attention(out)
        out = self.conv2(out)
        out = self.norm2(out)
        if self.apply_residual:
            out = out + identity
        return out


class MNASNetBackbone(nn.Module, ABC):

    def __init__(self, in_channels=3, out_channels=1280, stage_setting=None, width_multiplier=1.0, round_nearest=8, with_attention=True, reduction=4, attention_type='SqueezeAndExcitationBlock2D', conv_layer=None, norm_layer=None, act_layer=None):
        """
        MnasNet-A1实现
        参考Torchvision mnasnet实现，在Backbone之后添加一个Conv1x1-BN-ReLU，最终输出特征维度为1280
        :param in_channels: 输入通道数
        :param out_channels: 输出通道数
        :param stage_setting: 参数设置
        :param width_multiplier: 宽度乘法器
        :param round_nearest: 设置每一层通道数均为8的倍数
        :param with_attention: 是否使用注意力模块
        :param reduction: 衰减率
        :param attention_type: 注意力模块类型
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        """
        super(MNASNetBackbone, self).__init__()
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        block_layer = MNASNetUint
        base_planes = 32
        if stage_setting is None:
            stage_setting = [[3, 1, 1, 0, 1, 16], [3, 2, 6, 0, 2, 24], [5, 2, 3, 1, 3, 40], [3, 2, 6, 0, 4, 80], [3, 1, 6, 1, 2, 112], [5, 2, 6, 1, 3, 160], [3, 1, 6, 0, 1, 320]]
        base_planes = round_to_multiple_of(base_planes * width_multiplier, round_nearest)
        for i in range(len(stage_setting)):
            stage_setting[i][-1] = round_to_multiple_of(stage_setting[i][-1] * width_multiplier, round_nearest)
        self.first_stem, self.last_stem = make_stem(in_channels, base_planes, stage_setting[-1][-1], out_channels, conv_layer, norm_layer, act_layer)
        in_channels = base_planes
        for i, (kernel_size, stride, expansion_rate, with_attention_2, repeated, out_channels) in enumerate(stage_setting):
            features = list()
            for j in range(repeated):
                stride = stride if j == 0 else 1
                features.append(block_layer(in_channels, out_channels, stride=stride, kernel_size=kernel_size, expansion_rate=expansion_rate, with_attention=with_attention and with_attention_2, reduction=reduction, attention_type=attention_type, conv_layer=conv_layer, norm_layer=norm_layer, act_layer=act_layer))
                in_channels = out_channels
            stage_name = f'stage{i + 1}'
            self.add_module(stage_name, nn.Sequential(*features))
        self.stage_num = len(stage_setting)
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.kaiming_uniform_(m.weight, mode='fan_out', nonlinearity='sigmoid')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def forward(self, x):
        x = self.first_stem(x)
        for i in range(self.stage_num):
            stage = self.__getattr__(f'stage{i + 1}')
            x = stage(x)
        x = self.last_stem(x)
        return x


class MobileNetV1Block(nn.Module, ABC):

    def __init__(self, in_channels, out_channels, stride=1, padding=1, conv_layer=None, norm_layer=None, act_layer=None) ->None:
        """
        Block = depth-wise convolution + point-wise convolution
        :param in_channels: 输入通道数
        :param out_channels: 输出通道数
        :param stride: 卷积层步长
        :param padding: 卷积层零填充
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        """
        super(MobileNetV1Block, self).__init__()
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        self.conv1 = conv_layer(in_channels, in_channels, kernel_size=3, stride=stride, padding=padding, groups=in_channels, bias=False)
        self.bn1 = norm_layer(in_channels)
        self.relu = act_layer(inplace=True)
        self.conv2 = conv_layer(in_channels, out_channels, kernel_size=1, stride=1, bias=False)
        self.bn2 = norm_layer(out_channels)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        return x


def init_weights(modules):
    for m in modules:
        if isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
            nn.init.constant_(m.weight, 1)
            nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, 0, 0.01)
            if m.bias is not None:
                nn.init.zeros_(m.bias)


class MobileNetV1Backbone(nn.Module, ABC):

    def __init__(self, in_channels=3, base_channels=32, layer_channels=(64, 128, 128, 256, 256, 512, 512, 512, 512, 512, 512, 1024, 1024), strides=(1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 2), conv_layer=None, norm_layer=None, act_layer=None):
        """
        :param in_channels: 输入通道数
        :param base_channels: 第一个卷积层通道数
        :param layer_channels: 后续深度卷积层通道数
        :param strides: 卷积步长
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        """
        super(MobileNetV1Backbone, self).__init__()
        assert len(strides) == len(layer_channels)
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        self.stem = make_stem(in_channels, base_channels, conv_layer, norm_layer, act_layer)
        self.layer_num = self._make_dw(base_channels, layer_channels, strides, conv_layer, norm_layer, act_layer)
        init_weights(self.modules())

    def _make_dw(self, base_planes, layer_planes, strides, conv_layer, norm_layer, act_layer):
        layer_num = len(layer_planes)
        in_planes = base_planes
        padding = 1
        for i, (layer_plane, stride) in enumerate(zip(layer_planes, strides)):
            if i == layer_num - 1:
                padding = 4
            dw_layer = MobileNetV1Block(int(in_planes), int(layer_plane), stride=stride, padding=padding, conv_layer=conv_layer, norm_layer=norm_layer, act_layer=act_layer)
            in_planes = layer_plane
            layer_name = f'layer{i + 1}'
            self.add_module(layer_name, dw_layer)
        return layer_num

    def _forward_impl(self, x):
        x = self.stem(x)
        for i in range(self.layer_num):
            layer = self.__getattr__(f'layer{i + 1}')
            x = layer(x)
        return x

    def forward(self, x):
        return self._forward_impl(x)


class MobileNetV2InvertedResidual(nn.Module, ABC):

    def __init__(self, in_channels, out_channels, expansion_rate=1, stride=1, padding=1, conv_layer=None, norm_layer=None, act_layer=None):
        """
        MobileNetV2的反向残差块由一个膨胀卷积和一个深度可分离卷积(depth-wise conv + point-wise conv)组成
        参考torchvision实现:
        1. 当膨胀率大于1时，执行膨胀卷积操作；
        2. 当深度卷积步长为1且输入/输出通道数相同时，执行残差连接
        3. 反向残差块的最后不执行激活操作
        :param in_channels: 输入通道数
        :param out_channels: 输出通道数
        :param expansion_rate: 膨胀因子
        :param stride: 卷积层步长
        :param padding: 卷积层零填充
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        """
        super(MobileNetV2InvertedResidual, self).__init__()
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU6
        hidden_channels = int(expansion_rate * in_channels)
        features = list()
        if expansion_rate != 1:
            features.append(nn.Sequential(conv_layer(in_channels, hidden_channels, kernel_size=1, stride=1, bias=False), norm_layer(hidden_channels), act_layer(inplace=True)))
        features.append(nn.Sequential(conv_layer(hidden_channels, hidden_channels, kernel_size=3, stride=stride, padding=padding, bias=False, groups=hidden_channels), norm_layer(hidden_channels), act_layer(inplace=True)))
        features.append(nn.Sequential(conv_layer(hidden_channels, out_channels, kernel_size=1, stride=1, bias=False), norm_layer(out_channels)))
        self.conv = nn.Sequential(*features)
        self.use_res_connect = stride == 1 and in_channels == out_channels

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)


class MobileNetV2Block(nn.Module, ABC):

    def __init__(self, in_channels, out_channels, expansion_rate=1, repeat=1, stride=1, padding=1, conv_layer=None, norm_layer=None, act_layer=None):
        """
        repeat InvertedResidualUnit. if stride>1, make downsample to the first Unit
        :param in_channels: 输入通道数
        :param out_channels: 输出通道数
        :param expansion_rate: 膨胀因子
        :param repeat: 重复次数
        :param stride: 卷积层步长
        :param padding: 卷积层零填充
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        """
        super(MobileNetV2Block, self).__init__()
        features = list()
        for i in range(repeat):
            if i != 0:
                stride = 1
            features.append(MobileNetV2InvertedResidual(in_channels, out_channels, expansion_rate=expansion_rate, stride=stride, padding=padding, conv_layer=conv_layer, norm_layer=norm_layer, act_layer=act_layer))
            in_channels = out_channels
        self.conv = nn.Sequential(*features)

    def forward(self, x):
        return self.conv(x)


class MobileNetV2Backbone(nn.Module, ABC):

    def __init__(self, in_channels=3, out_channels=1280, base_channels=32, stride=2, padding=1, inverted_residual_setting=None, conv_layer=None, norm_layer=None, act_layer=None):
        """
        :param in_channels: 输入通道数
        :param out_channels: 输出通道数
        :param base_channels: 第一个卷积层通道数
        :param stride: 第一个卷积层步长
        :param padding: 第一个卷积层零填充
        :param inverted_residual_setting: 反向残差块设置
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        """
        super(MobileNetV2Backbone, self).__init__()
        if inverted_residual_setting is None:
            inverted_residual_setting = [[1, 16, 1, 1], [6, 24, 2, 2], [6, 32, 3, 2], [6, 64, 4, 2], [6, 96, 3, 1], [6, 160, 3, 2], [6, 320, 1, 1]]
        else:
            assert len(inverted_residual_setting[0]) == 4
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU6
        kernel_size = 3
        stride = 2
        self.first_stem = make_stem(in_channels, base_channels, kernel_size, stride, conv_layer, norm_layer, act_layer)
        in_channels = base_channels
        for i, (t, c, n, s) in enumerate(inverted_residual_setting):
            feature_dims = c
            layer = MobileNetV2Block(in_channels, feature_dims, expansion_rate=t, repeat=n, stride=s, conv_layer=conv_layer, norm_layer=norm_layer, act_layer=act_layer)
            layer_name = f'layer{i + 1}'
            self.add_module(layer_name, layer)
            in_channels = feature_dims
        kernel_size = 1
        stride = 1
        self.last_stem = make_stem(in_channels, out_channels, kernel_size, stride, conv_layer, norm_layer, act_layer)
        self.layer_num = len(inverted_residual_setting)
        init_weights(self.modules())

    def forward(self, x):
        x = self.first_stem(x)
        for i in range(self.layer_num):
            layer = self.__getattr__(f'layer{i + 1}')
            x = layer(x)
        x = self.last_stem(x)
        return x


class MobileNetV3Unit(nn.Module, ABC):

    def __init__(self, in_channels, inner_channels, out_channels, stride=1, kernel_size=3, round_nearest=8, with_attention=True, reduction=4, attention_type='SqueezeAndExcitationBlock2D', conv_layer=None, norm_layer=None, act_layer=None, sigmoid_type=None):
        """
        MobileNetV2的反向残差构建块 + Squeeze-and-Excite + h-swish
        :param in_channels: 输入通道数
        :param inner_channels: 中间膨胀块大小
        :param out_channels: 输出通道数
        :param stride: 步长
        :param kernel_size: 卷积核大小
        :param round_nearest: 设置每一层通道数均为8的倍数
        :param with_attention: 是否使用注意力模块
        :param reduction: 衰减率
        :param attention_type: 注意力模块类型
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        :param sigmoid_type: sigmoid类型
        """
        super(MobileNetV3Unit, self).__init__()
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.Hardswish
        if sigmoid_type is None:
            sigmoid_type = 'HSigmoid'
        self.with_attention = with_attention
        self.expansion = nn.Sequential(conv_layer(in_channels, inner_channels, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), bias=False), norm_layer(inner_channels), act_layer(inplace=True)) if in_channels != inner_channels else nn.Identity()
        if kernel_size == 3:
            padding = 1
        elif kernel_size == 5:
            padding = 2
        else:
            padding = 0
        self.conv1 = conv_layer(inner_channels, inner_channels, kernel_size=(kernel_size, kernel_size), stride=(stride, stride), padding=(padding, padding), bias=False, groups=inner_channels)
        self.norm1 = norm_layer(inner_channels)
        self.act = act_layer(inplace=True)
        if self.with_attention:
            self.attention = make_attention_block(inner_channels, reduction, attention_type, sigmoid_type=sigmoid_type, is_round=True, round_nearest=round_nearest, bias=True)
        self.conv2 = conv_layer(inner_channels, out_channels, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), bias=False)
        self.norm2 = norm_layer(out_channels)
        self.apply_residual = in_channels == out_channels and stride == 1

    def forward(self, x):
        identity = x
        out = self.expansion(x)
        out = self.conv1(out)
        out = self.norm1(out)
        out = self.act(out)
        if self.with_attention:
            out = self.attention(out)
        out = self.conv2(out)
        out = self.norm2(out)
        if self.apply_residual:
            out = out + identity
        return out


def relu_or_hswish(name):
    if name == 'RE':
        return nn.ReLU
    elif name == 'HS':
        return nn.Hardswish
    else:
        raise IOError(f'{name} does not exist')


class MobileNetV3Backbone(nn.Module, ABC):

    def __init__(self, in_channels=3, base_channels=16, out_channels=960, layer_setting=None, width_multiplier=1.0, round_nearest=8, with_attention=True, reduction=4, attention_type='SqueezeAndExcitationBlock2D', conv_layer=None, norm_layer=None, act_layer=None, sigmoid_type=None):
        """
        :param in_channels: 输入通道数
        :param base_channels: 基础通道数
        :param out_channels: 输出通道数
        :param width_multiplier: 宽度乘法器
        :param round_nearest: 设置每一层通道数均为8的倍数
        :param with_attention: 是否使用注意力模块
        :param reduction: 衰减率
        :param attention_type: 注意力模块类型
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        :param sigmoid_type: sigmoid类型
        """
        super(MobileNetV3Backbone, self).__init__()
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.Hardswish
        if sigmoid_type is None:
            sigmoid_type = 'HSigmoid'
        block_layer = MobileNetV3Unit
        if layer_setting is None:
            layer_setting = [[3, 1, 16, 0, 'RE', 16], [3, 2, 64, 0, 'RE', 24], [3, 1, 72, 0, 'RE', 24], [5, 2, 72, 1, 'RE', 40], [5, 1, 120, 1, 'RE', 40], [5, 1, 120, 1, 'RE', 40], [3, 2, 240, 0, 'HS', 80], [3, 1, 200, 0, 'HS', 80], [3, 1, 184, 0, 'HS', 80], [3, 1, 184, 0, 'HS', 80], [3, 1, 480, 1, 'HS', 112], [3, 1, 672, 1, 'HS', 112], [5, 2, 672, 1, 'HS', 160], [5, 1, 960, 1, 'HS', 160], [5, 1, 960, 1, 'HS', 160]]
        base_channels = round_to_multiple_of(base_channels * width_multiplier, round_nearest)
        out_channels = round_to_multiple_of(out_channels * width_multiplier, round_nearest)
        for i in range(len(layer_setting)):
            layer_setting[i][2] = round_to_multiple_of(layer_setting[i][2], round_nearest)
            layer_setting[i][-1] = round_to_multiple_of(layer_setting[i][-1] * width_multiplier, round_nearest)
        self.first_stem, self.last_stem = make_stem(in_channels, base_channels, layer_setting[-1][-1], out_channels, conv_layer, norm_layer, act_layer)
        features = list()
        in_channels = base_channels
        for i, (kernel_size, stride, inner_planes, with_attention_2, non_linearity, out_channels) in enumerate(layer_setting):
            act_layer = relu_or_hswish(non_linearity)
            features.append(block_layer(in_channels, inner_planes, out_channels, stride=stride, kernel_size=kernel_size, with_attention=with_attention_2 and with_attention, reduction=reduction, attention_type=attention_type, conv_layer=conv_layer, norm_layer=norm_layer, act_layer=act_layer, sigmoid_type=sigmoid_type))
            in_channels = out_channels
        self.add_module('features', nn.Sequential(*features))
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.kaiming_uniform_(m.weight, mode='fan_out', nonlinearity='sigmoid')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def forward(self, x):
        x = self.first_stem(x)
        x = self.features(x)
        x = self.last_stem(x)
        return x


class BasicBlock(nn.Module, ABC):
    expansion = 1

    def __init__(self, in_channels, out_channels, stride=1, downsample=None, groups=1, base_width=64, with_attention=False, reduction=16, attention_type='SqueezeAndExcitationBlock2D', conv_layer=None, norm_layer=None, act_layer=None, **kwargs):
        """
        使用两个3x3卷积，如果进行下采样，那么使用第一个卷积层对输入空间尺寸进行减半操作
        参考Torchvision实现
        对于注意力模块，有两种嵌入方式：
        1. 对于Squeeze-And-Excitation或者Global Context操作，在残差连接中（after 1x1）嵌入；
        2. 对于NonLocal或者SimplifiedNonLoal，在Block完成计算后（after add）嵌入。
        :param in_channels:　输入通道数
        :param out_channels:　输出通道数
        :param stride:　步长
        :param downsample:　下采样
        :param groups:　cardinality
        :param base_width:　基础宽度
        :param with_attention:　是否使用注意力模块
        :param reduction:　衰减率
        :param attention_type:　注意力模块类型
        :param conv_layer:　卷积层类型
        :param norm_layer:　归一化层类型
        :param act_layer:　激活层类型
        :param kwargs:　其他参数
        """
        super(BasicBlock, self).__init__()
        assert with_attention in (0, 1)
        assert attention_type in ['GlobalContextBlock2D', 'SimplifiedNonLocal2DEmbeddedGaussian', 'NonLocal2DEmbeddedGaussian', 'SqueezeAndExcitationBlock2D']
        if groups != 1 or base_width != 64:
            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        self.down_sample = downsample
        self.conv1 = conv_layer(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = norm_layer(out_channels)
        self.conv2 = conv_layer(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = norm_layer(out_channels)
        self.relu = act_layer(inplace=True)
        self.attention_after_1x1 = None
        self.attention_after_add = None
        if with_attention and attention_type in ['SqueezeAndExcitationBlock2D', 'GlobalContextBlock2D']:
            self.attention_after_1x1 = make_attention_block(out_channels * self.expansion, reduction, attention_type)
            self.attention_after_add = None
        if with_attention and attention_type in ['NonLocal2DEmbeddedGaussian', 'SimplifiedNonLocal2DEmbeddedGaussian']:
            self.attention_after_1x1 = None
            self.attention_after_add = make_attention_block(out_channels * self.expansion, reduction, attention_type)

    def forward(self, x):
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if self.attention_after_1x1 is not None:
            out = self.attention_after_1x1(out)
        if self.down_sample is not None:
            identity = self.down_sample(x)
        out += identity
        out = self.relu(out)
        if self.attention_after_add is not None:
            out = self.attention_after_add(out)
        return out


class Bottleneck(nn.Module, ABC):
    expansion = 4

    def __init__(self, in_planes, out_planes, stride=1, down_sample=None, groups=1, base_width=64, with_attention=False, reduction=16, attention_type='SqueezeAndExcitationBlock2D', conv_layer=None, norm_layer=None, act_layer=None, use_avg=False, fast_avg=False, **kwargs):
        """
        依次执行大小为1x1、3x3、1x1的卷积操作，如果进行下采样，那么使用第二个卷积层对输入空间尺寸进行减半操作
        参考Torchvision实现
        对于注意力模块，有两种嵌入方式：
        1. 对于Squeeze-And-Excitation或者Global Context操作，在残差连接中（after 1x1）嵌入；
        2. 对于NonLocal或者SimplifiedNonLoal，在Block完成计算后（after add）嵌入。
        参考ResNeSt实现，使用AvgPool替代Conv进行下采样操作，当前仅对Bottleneck进行实现
        :param in_channels:　输入通道数
        :param out_channels:　输出通道数
        :param stride:　步长
        :param downsample:　下采样
        :param groups:　cardinality
        :param base_width:　基础宽度
        :param with_attention:　是否使用注意力模块
        :param reduction:　衰减率
        :param attention_type:　注意力模块类型
        :param conv_layer:　卷积层类型
        :param norm_layer:　归一化层类型
        :param act_layer:　激活层类型
        :param kwargs:　其他参数
        """
        super(Bottleneck, self).__init__()
        assert with_attention in (0, 1)
        assert attention_type in ['GlobalContextBlock2D', 'SimplifiedNonLocal2DEmbeddedGaussian', 'NonLocal2DEmbeddedGaussian', 'SqueezeAndExcitationBlock2D']
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        self.down_sample = down_sample
        width = int(out_planes * (base_width / 64.0)) * groups
        self.conv1 = conv_layer(in_planes, width, kernel_size=1, stride=1, bias=False)
        self.bn1 = norm_layer(width)
        self.fast_avg = fast_avg
        self.avg = None
        if use_avg and stride > 1:
            self.avg = nn.AvgPool2d(kernel_size=3, stride=stride, padding=1)
            stride = 1
        self.conv2 = conv_layer(width, width, kernel_size=3, stride=stride, padding=1, bias=False, groups=groups)
        self.bn2 = norm_layer(width)
        self.conv3 = conv_layer(width, out_planes * self.expansion, kernel_size=1, stride=1, bias=False)
        self.bn3 = norm_layer(out_planes * self.expansion)
        self.relu = act_layer(inplace=True)
        self.attention_after_1x1 = None
        self.attention_after_add = None
        if with_attention and attention_type in ['SqueezeAndExcitationBlock2D', 'GlobalContextBlock2D']:
            self.attention_after_1x1 = make_attention_block(out_planes * self.expansion, reduction, attention_type)
            self.attention_after_add = None
        if with_attention and attention_type in ['NonLocal2DEmbeddedGaussian', 'SimplifiedNonLocal2DEmbeddedGaussian']:
            self.attention_after_1x1 = None
            self.attention_after_add = make_attention_block(out_planes * self.expansion, reduction, attention_type)

    def forward(self, x):
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        if self.avg is not None and self.fast_avg:
            out = self.avg(out)
        out = self.conv2(out)
        if self.avg is not None and not self.fast_avg:
            out = self.avg(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv3(out)
        out = self.bn3(out)
        if self.attention_after_1x1 is not None:
            out = self.attention_after_1x1(out)
        if self.down_sample is not None:
            identity = self.down_sample(x)
        out += identity
        out = self.relu(out)
        if self.attention_after_add is not None:
            out = self.attention_after_add(out)
        return out


class rSoftMax(nn.Module):

    def __init__(self, radix, cardinality):
        super().__init__()
        self.radix = radix
        self.cardinality = cardinality

    def forward(self, x):
        batch = x.size(0)
        if self.radix > 1:
            x = x.view(batch, self.cardinality, self.radix, -1).transpose(1, 2)
            x = F.softmax(x, dim=1)
            x = x.reshape(batch, -1)
        else:
            x = torch.sigmoid(x)
        return x


class SplitAttentionConv2d(nn.Module, ABC):

    def __init__(self, in_channels, out_channels, radix=2, groups=1, reduction_rate=4, default_channels: int=32, dimension: int=2):
        """
        Implementation of SplitAttention in ResNetSt, refer to
        1. https://github.com/open-mmlab/mmdetection/blob/master/mmdet/models/backbones/resnest.py
        2. https://github.com/zhanghang1989/ResNeSt/blob/73b43ba63d1034dbf3e96b3010a8f2eb4cc3854f/resnest/torch/splat.py
        Partial reference ./selective_kernel_conv2d.py implementation
        :param in_channels:
        :param out_channels:
        :param radix:
        :param groups:
        :param reduction_rate:
        :param default_channels:
        :param dimension:
        """
        super(SplitAttentionConv2d, self).__init__()
        self.split = nn.Sequential(nn.Conv2d(in_channels, out_channels * radix, kernel_size=3, stride=1, padding=1, bias=False, groups=groups * radix), nn.BatchNorm2d(out_channels * radix), nn.ReLU(inplace=True))
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        inner_channels = max(in_channels * radix // reduction_rate, 32)
        self.compact = nn.Sequential(nn.Conv2d(out_channels, inner_channels, kernel_size=1, stride=1, padding=0, bias=True, groups=groups), nn.BatchNorm2d(inner_channels), nn.ReLU(inplace=True))
        self.select = nn.Conv2d(inner_channels, out_channels * radix, kernel_size=1, stride=1, groups=groups)
        self.rsoftmax = rSoftMax(radix, groups)
        self.dimension = dimension
        self.out_channels = out_channels
        self.radix = radix
        self.groups = groups
        init_weights(self.modules())

    def forward(self, x):
        N, C, H, W = x.shape[:4]
        out = self.split(x)
        split_out = torch.stack(torch.split(out, self.out_channels, dim=1))
        u = torch.sum(split_out, dim=0)
        s = self.pool(u)
        z = self.compact(s)
        c = self.select(z)
        softmax_c = self.rsoftmax(c).view(N, -1, 1, 1)
        attens = torch.split(softmax_c, self.out_channels, dim=1)
        v = sum([(att * split) for att, split in zip(attens, split_out)])
        return v.contiguous()


class ResNeStBlock(nn.Module, ABC):
    expansion = 4

    def __init__(self, in_channels, out_channels, stride=1, downsample=None, groups=1, base_width=64, with_attention=False, reduction=4, attention_type='SqueezeAndExcitationBlock2D', conv_layer=None, norm_layer=None, act_layer=None, radix=1, fast_avg=False, **kwargs):
        """
        依次执行大小为1x1、3x3、1x1的卷积操作，如果进行下采样，那么使用第二个卷积层对输入空间尺寸进行减半操作
        参考Torchvision实现
        对于注意力模块，有两种嵌入方式：
        1. 对于Squeeze-And-Excitation或者Global Context操作，在残差连接中（after 1x1）嵌入；
        2. 对于NonLocal或者SimplifiedNonLoal，在Block完成计算后（after add）嵌入。
        对于Selective Kernel Conv2d，替换3x3卷积层；
        对于下采样操作，参考
        ResNeSt-fast setting, the effective average downsampling is applied prior to the
        3 × 3 convolution to avoid introducing extra computational costs in the model.
        With the downsampling operation moved after the convolutional layer, ResNeSt-
        50 achieves 81.13% accuracy
        在3x3卷积层之前（fast设置）或者之后执行AvgPool2d操作
        :param in_channels: 输入通道数
        :param out_channels: 输出通道数
        :param stride: 步长
        :param downsample: 下采样
        :param groups: cardinality
        :param base_width: 基础宽度
        :param with_attention: 是否使用注意力模块
        :param reduction: 衰减率
        :param attention_type: 注意力模块类型
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        :param radix: 每个group中的分离数
        :param fast_avg: 在3x3之前执行下采样操作
        :param kwargs: 其他参数
        """
        super(ResNeStBlock, self).__init__()
        assert radix > 0
        assert with_attention in (0, 1)
        assert attention_type in ['GlobalContextBlock2D', 'SimplifiedNonLocal2DEmbeddedGaussian', 'NonLocal2DEmbeddedGaussian', 'SqueezeAndExcitationBlock2D']
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        self.radix = radix
        self.down_sample = downsample
        width = int(out_channels * (base_width / 64.0)) * groups
        self.conv1 = conv_layer(in_channels, width, kernel_size=1, stride=1, bias=False)
        self.bn1 = norm_layer(width)
        self.conv2 = SplitAttentionConv2d(width, width, radix=radix, groups=groups, reduction_rate=reduction)
        if self.radix == 0:
            self.bn2 = norm_layer(width)
        self.conv3 = conv_layer(width, out_channels * self.expansion, kernel_size=1, stride=1, bias=False)
        self.bn3 = norm_layer(out_channels * self.expansion)
        self.relu = act_layer(inplace=True)
        self.attention_after_1x1 = None
        self.attention_after_add = None
        if with_attention and attention_type in ['SqueezeAndExcitationBlock2D', 'GlobalContextBlock2D']:
            self.attention_after_1x1 = make_attention_block(out_channels * self.expansion, reduction, attention_type)
            self.attention_after_add = None
        if with_attention and attention_type in ['NonLocal2DEmbeddedGaussian', 'SimplifiedNonLocal2DEmbeddedGaussian']:
            self.attention_after_1x1 = None
            self.attention_after_add = make_attention_block(out_channels * self.expansion, reduction, attention_type)
        self.fast_avg = fast_avg
        self.avg = None
        if stride > 1:
            self.avg = nn.AvgPool2d(kernel_size=3, stride=stride, padding=1)

    def forward(self, x):
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        if self.fast_avg and self.avg is not None:
            out = self.avg(out)
        out = self.conv2(out)
        if self.radix == 0:
            out = self.bn2(out)
            out = self.relu(out)
        if not self.fast_avg and self.avg is not None:
            out = self.avg(out)
        out = self.conv3(out)
        out = self.bn3(out)
        if self.attention_after_1x1 is not None:
            out = self.attention_after_1x1(out)
        if self.down_sample is not None:
            identity = self.down_sample(x)
        out += identity
        out = self.relu(out)
        if self.attention_after_add is not None:
            out = self.attention_after_add(out)
        return out


class PlaceHolder(nn.Module, ABC):
    """
    @deprecated. pytorch is implemented, using nn.Identity()
    """

    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x


class ResNet3DBasicBlock(nn.Module, ABC):
    expansion = 1

    def __init__(self, in_planes, out_planes, spatial_stride=1, temporal_stride=1, down_sample=None, inflate=False, inflate_style='3x1x1', groups=1, base_width=64, conv_layer=None, norm_layer=None, act_layer=None):
        """
        使用两个Tx3x3卷积，如果进行下采样，那么使用第一个卷积层对输入时空尺寸进行减半操作
        如果执行膨胀操作，仅作用于第一个卷积层，第二个卷积层的kernel_size大小为1x3x3
        :param in_planes: 输入通道数
        :param out_planes: 输出通道数
        :param spatial_stride: 空间步长
        :param temporal_stride: 时间步长
        :param down_sample: 下采样
        :param inflate: 是否膨胀
        :param inflate_style: 膨胀类型，作用于Bottleneck
        :param groups: cardinality
        :param base_width: 基础宽度
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        """
        super(ResNet3DBasicBlock, self).__init__()
        assert inflate_style in ('3x1x1', '3x3x3')
        if groups != 1 or base_width != 64:
            raise ValueError('ResNet3DBasicBlock only supports groups=1 and base_width=64')
        if conv_layer is None:
            conv_layer = nn.Conv3d
        if norm_layer is None:
            norm_layer = nn.BatchNorm3d
        if act_layer is None:
            act_layer = nn.ReLU
        if inflate:
            conv1_kernel_size = 3, 3, 3
            conv1_stride = temporal_stride, spatial_stride, spatial_stride
            conv1_padding = 1, 1, 1
        else:
            conv1_kernel_size = 1, 3, 3
            conv1_stride = 1, spatial_stride, spatial_stride
            conv1_padding = 0, 1, 1
        self.down_sample = down_sample
        self.conv1 = conv_layer(in_planes, out_planes, kernel_size=conv1_kernel_size, stride=conv1_stride, padding=conv1_padding, bias=False)
        self.bn1 = norm_layer(out_planes)
        self.conv2 = conv_layer(out_planes, out_planes, kernel_size=(1, 3, 3), stride=(1, 1, 1), padding=(0, 1, 1), bias=False)
        self.bn2 = norm_layer(out_planes)
        self.relu = act_layer(inplace=True)

    def forward(self, x):
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if self.down_sample is not None:
            identity = self.down_sample(x)
        out += identity
        out = self.relu(out)
        return out


class ResNet3DBottleneck(nn.Module, ABC):
    expansion = 4

    def __init__(self, in_planes, out_planes, spatial_stride=1, temporal_stride=1, down_sample=None, inflate=False, inflate_style='3x1x1', groups=1, base_width=64, conv_layer=None, norm_layer=None, act_layer=None):
        """
        依次执行大小为Tx1x1、Tx3x3、1x1x1的卷积操作，
        如果进行下采样，那么使用第二个卷积层对输入空间尺寸进行减半操作
        分两种膨胀类型：3x1x1和3x3x3
        :param in_planes: 输入通道数
        :param out_planes: 输出通道数
        :param spatial_stride: 空间步长
        :param temporal_stride: 时间步长
        :param down_sample: 下采样
        :param inflate: 是否膨胀
        :param inflate_style: 膨胀类型，作用于Bottleneck
        :param groups: cardinality
        :param base_width: 基础宽度
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        """
        super(ResNet3DBottleneck, self).__init__()
        assert inflate_style in ('3x1x1', '3x3x3')
        if conv_layer is None:
            conv_layer = nn.Conv3d
        if norm_layer is None:
            norm_layer = nn.BatchNorm3d
        if act_layer is None:
            act_layer = nn.ReLU
        if inflate:
            if inflate_style == '3x1x1':
                conv1_kernel_size = 3, 1, 1
                conv1_stride = temporal_stride, 1, 1
                conv1_padding = 1, 0, 0
                conv2_kernel_size = 1, 3, 3
                conv2_stride = 1, spatial_stride, spatial_stride
                conv2_padding = 0, 1, 1
            else:
                conv1_kernel_size = 1, 1, 1
                conv1_stride = 1, 1, 1
                conv1_padding = 0, 0, 0
                conv2_kernel_size = 3, 3, 3
                conv2_stride = temporal_stride, spatial_stride, spatial_stride
                conv2_padding = 1, 1, 1
        else:
            conv1_kernel_size = 1, 1, 1
            conv1_stride = 1, 1, 1
            conv1_padding = 0, 0, 0
            conv2_kernel_size = 1, 3, 3
            conv2_stride = 1, spatial_stride, spatial_stride
            conv2_padding = 0, 1, 1
        self.down_sample = down_sample
        width = int(out_planes * (base_width / 64.0)) * groups
        self.conv1 = conv_layer(in_planes, width, kernel_size=conv1_kernel_size, stride=conv1_stride, padding=conv1_padding, bias=False)
        self.bn1 = norm_layer(width)
        self.conv2 = conv_layer(width, width, kernel_size=conv2_kernel_size, stride=conv2_stride, padding=conv2_padding, bias=False, groups=groups)
        self.bn2 = norm_layer(width)
        self.conv3 = conv_layer(width, out_planes * self.expansion, kernel_size=(1, 1, 1), stride=(1, 1, 1), bias=False)
        self.bn3 = norm_layer(out_planes * self.expansion)
        self.relu = act_layer(inplace=True)

    def forward(self, x):
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv3(out)
        out = self.bn3(out)
        if self.down_sample is not None:
            identity = self.down_sample(x)
        out += identity
        out = self.relu(out)
        return out


class ResNet3DBackbone(nn.Module, ABC):

    def __init__(self, in_planes=3, base_planes=64, conv1_kernel=(1, 7, 7), conv1_stride=(1, 2, 2), conv1_padding=(0, 3, 3), pool1_kernel=(1, 3, 3), pool1_stride=(1, 2, 2), pool1_padding=(0, 1, 1), with_pool2=False, layer_planes=(64, 128, 256, 512), layer_blocks=(2, 2, 2, 2), down_samples=(0, 1, 1, 1), temporal_strides=(1, 1, 1, 1), inflate_list=(0, 0, 0, 0), inflate_style='3x1x1', groups=1, width_per_group=64, block_layer=None, conv_layer=None, norm_layer=None, act_layer=None, zero_init_residual=False, state_dict_2d=None):
        super(ResNet3DBackbone, self).__init__()
        if block_layer is None:
            block_layer = ResNet3DBasicBlock
        if conv_layer is None:
            conv_layer = nn.Conv3d
        if norm_layer is None:
            norm_layer = nn.BatchNorm3d
        if act_layer is None:
            act_layer = nn.ReLU
        self._make_stem(in_planes, base_planes, conv1_kernel, conv1_stride, conv1_padding, pool1_kernel, pool1_stride, pool1_padding, with_pool2, conv_layer, norm_layer, act_layer)
        inplanes = base_planes
        for i in range(len(layer_blocks)):
            res_layer = self._make_res_layer(inplanes, layer_planes[i], layer_blocks[i], down_samples[i], temporal_strides[i], inflate_list[i], inflate_style, groups, width_per_group, block_layer, conv_layer, norm_layer, act_layer)
            inplanes = layer_planes[i] * block_layer.expansion
            layer_name = f'layer{i + 1}'
            self.add_module(layer_name, res_layer)
        self.init_weights(zero_init_residual, state_dict_2d)

    def _make_stem(self, in_planes, base_planes, conv1_kernel, conv1_stride, conv1_padding, pool1_kernel, pool1_stride, pool1_padding, with_pool2, conv_layer, norm_layer, act_layer):
        self.conv1 = conv_layer(in_planes, base_planes, kernel_size=conv1_kernel, stride=conv1_stride, padding=conv1_padding, bias=False)
        self.bn1 = norm_layer(base_planes)
        self.relu = act_layer(inplace=True)
        self.pool = nn.MaxPool3d(kernel_size=pool1_kernel, stride=pool1_stride, padding=pool1_padding)
        if with_pool2:
            self.pool2 = nn.MaxPool3d(kernel_size=(3, 1, 1), stride=(2, 1, 1), padding=(1, 0, 0))
        else:
            self.pool2 = PlaceHolder()

    def _make_res_layer(self, in_planes, out_planes, block_num, with_down_sample, temporal_stride, inflate, inflate_style, groups, width_per_group, block_layer, conv_layer, norm_layer, act_layer):
        inflates = inflate if not isinstance(inflate, int) else (inflate,) * block_num
        assert len(inflates) == block_num
        assert inflate_style in ('3x1x1', '3x3x3')
        spatial_stride = 2 if with_down_sample else 1
        expansion = block_layer.expansion
        if with_down_sample or in_planes != out_planes * expansion:
            down_sample = nn.Sequential(conv_layer(in_planes, out_planes * expansion, kernel_size=1, stride=(temporal_stride, spatial_stride, spatial_stride), bias=False), norm_layer(out_planes * expansion))
        else:
            down_sample = None
        blocks = list()
        blocks.append(block_layer(in_planes, out_planes, spatial_stride, temporal_stride, down_sample, inflates[0], inflate_style, groups, width_per_group, conv_layer, norm_layer, act_layer))
        in_planes = out_planes * expansion
        spatial_stride = 1
        temporal_stride = 1
        down_sample = None
        for i in range(1, block_num):
            blocks.append(block_layer(in_planes, out_planes, spatial_stride, temporal_stride, down_sample, inflates[i], inflate_style, groups, width_per_group, conv_layer, norm_layer, act_layer))
        return nn.Sequential(*blocks)

    def init_weights(self, zero_init_residual, state_dict_2d):
        for m in self.modules():
            if isinstance(m, nn.Conv3d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm3d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, ResNet3DBottleneck):
                    nn.init.constant_(m.bn3.weight, 0)
                elif isinstance(m, ResNet3DBasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)
        if state_dict_2d:

            def _inflate_conv_params(conv3d, state_dict_2d, module_name_2d, inflated_param_names):
                """Inflate a conv module from 2d to 3d.

                Args:
                    conv3d (nn.Module): The destination conv3d module.
                    state_dict_2d (OrderedDict): The state dict of pretrained 2d model.
                    module_name_2d (str): The name of corresponding conv module in the
                        2d model.
                    inflated_param_names (list[str]): List of parameters that have been
                        inflated.
                """
                weight_2d_name = module_name_2d + '.weight'
                if weight_2d_name in state_dict_2d.keys():
                    conv2d_weight = state_dict_2d[weight_2d_name]
                    kernel_t = conv3d.weight.data.shape[2]
                    new_weight = conv2d_weight.data.unsqueeze(2).expand_as(conv3d.weight) / kernel_t
                    conv3d.weight.data.copy_(new_weight)
                    inflated_param_names.append(weight_2d_name)
                    if getattr(conv3d, 'bias') is not None:
                        bias_2d_name = module_name_2d + '.bias'
                        conv3d.bias.data.copy_(state_dict_2d[bias_2d_name])
                        inflated_param_names.append(bias_2d_name)

            def _inflate_bn_params(bn3d, state_dict_2d, module_name_2d, inflated_param_names):
                """Inflate a norm module from 2d to 3d.

                Args:
                    bn3d (nn.Module): The destination bn3d module.
                    state_dict_2d (OrderedDict): The state dict of pretrained 2d model.
                    module_name_2d (str): The name of corresponding bn module in the
                        2d model.
                    inflated_param_names (list[str]): List of parameters that have been
                        inflated.
                """
                for param_name, param in bn3d.named_parameters():
                    param_2d_name = f'{module_name_2d}.{param_name}'
                    if param_2d_name in state_dict_2d.keys():
                        param_2d = state_dict_2d[param_2d_name]
                        param.data.copy_(param_2d)
                        inflated_param_names.append(param_2d_name)
                for param_name, param in bn3d.named_buffers():
                    param_2d_name = f'{module_name_2d}.{param_name}'
                    if param_2d_name in state_dict_2d:
                        param_2d = state_dict_2d[param_2d_name]
                        param.data.copy_(param_2d)
                        inflated_param_names.append(param_2d_name)
            inflated_param_names = []
            for name, module in self.named_modules():
                if isinstance(module, nn.Conv3d):
                    _inflate_conv_params(module, state_dict_2d, name, inflated_param_names)
                if isinstance(module, (nn.BatchNorm3d, nn.GroupNorm)):
                    _inflate_bn_params(module, state_dict_2d, name, inflated_param_names)
            remaining_names = set(state_dict_2d.keys()) - set(inflated_param_names)
            if remaining_names:
                None

    def _forward_impl(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.pool(x)
        x = self.layer1(x)
        x = self.pool2(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        return x

    def forward(self, x):
        return self._forward_impl(x)


def make_res_layer(in_channels, out_channels, block_num, with_sample, groups, width_per_group, with_attention, reduction, attention_type, block_layer, conv_layer, norm_layer, act_layer, use_avg, fast_avg, **kwargs):
    """

    :param in_channels: 输入通道数
    :param out_channels: 输出通道数
    :param block_num: 块个数
    :param with_sample: 是否执行空间下采样
    :param groups: cardinality
    :param width_per_group: 每组的宽度
    :param with_attention: 是否使用注意力模块
    :param reduction: 衰减率
    :param attention_type: 注意力模块类型
    :param block_layer: 块类型
    :param conv_layer: 卷积层类型
    :param norm_layer: 归一化层类型
    :param act_layer: 激活层类型
    :param use_avg: 是否使用AvgPool进行下采样
    :param fast_avg: 在3x3之前执行下采样操作
    :param kwargs: 其他参数
    :return:
    """
    assert isinstance(with_attention, (int, tuple))
    assert with_attention in (0, 1) if isinstance(with_attention, int) else len(with_attention) == block_num
    with_attentions = with_attention if isinstance(with_attention, tuple) else [with_attention] * block_num
    stride = 2 if with_sample else 1
    expansion = block_layer.expansion
    if with_sample:
        down_sample = nn.Sequential(nn.AvgPool2d(kernel_size=2, stride=2), conv_layer(in_channels, out_channels * expansion, kernel_size=1, stride=1, bias=False), norm_layer(out_channels * expansion))
    elif in_channels != out_channels * expansion:
        down_sample = nn.Sequential(conv_layer(in_channels, out_channels * expansion, kernel_size=1, stride=1, bias=False), norm_layer(out_channels * expansion))
    else:
        down_sample = None
    blocks = list()
    blocks.append(block_layer(in_channels, out_channels, stride, down_sample, groups, width_per_group, with_attentions[0], reduction, attention_type, conv_layer, norm_layer, act_layer, use_avg=use_avg, fast_avg=fast_avg, **kwargs))
    in_channels = out_channels * expansion
    stride = 1
    down_sample = None
    for i in range(1, block_num):
        blocks.append(block_layer(in_channels, out_channels, stride, down_sample, groups, width_per_group, with_attentions[i], reduction, attention_type, conv_layer, norm_layer, act_layer, use_avg=use_avg, fast_avg=fast_avg, **kwargs))
    return nn.Sequential(*blocks)


class ResNetBackbone(nn.Module, ABC):

    def __init__(self, in_channels=3, base_channels=64, layer_channels=(64, 128, 256, 512), layer_blocks=(2, 2, 2, 2), down_samples=(0, 1, 1, 1), groups=1, width_per_group=64, with_attentions=(0, 0, 0, 0), reduction=16, attention_type='SqueezeAndExcitationBlock2D', block_layer=None, conv_layer=None, norm_layer=None, act_layer=None, zero_init_residual=False, **kwargs):
        """
        参考Torchvision实现，适用于torchvision预训练模型加载
        :param in_channels: 输入通道数
        :param base_channels: 基础通道数
        :param layer_channels: 每一层通道数
        :param layer_blocks: 每一层块个数
        :param down_samples: 是否执行空间下采样
        :param groups: cardinality
        :param width_per_group: 每组的宽度
        :param with_attentions: 是否使用注意力模块
        :param reduction: 衰减率
        :param attention_type: 注意力模块类型
        :param block_layer: 块类型
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        :param zero_init_residual: 零初始化残差连接
        :param kwargs: 其他参数
        """
        super(ResNetBackbone, self).__init__()
        assert len(layer_channels) == len(layer_blocks) == len(down_samples) == len(with_attentions)
        if block_layer is None:
            block_layer = BasicBlock
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        self.stem = make_stem(in_channels, base_channels, conv_layer, norm_layer, act_layer)
        in_channels = base_channels
        for i in range(len(layer_blocks)):
            res_layer = make_res_layer(in_channels, layer_channels[i], layer_blocks[i], down_samples[i], groups, width_per_group, with_attentions[i], reduction, attention_type, block_layer, conv_layer, norm_layer, act_layer, **kwargs)
            in_channels = layer_channels[i] * block_layer.expansion
            layer_name = f'layer{i + 1}'
            self.add_module(layer_name, res_layer)
        self.layer_num = len(layer_blocks)
        self.init_weights(zero_init_residual)

    def init_weights(self, zero_init_residual):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck):
                    nn.init.constant_(m.bn3.weight, 0)
                elif isinstance(m, BasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)

    def _forward_impl(self, x):
        x = self.stem(x)
        for i in range(self.layer_num):
            layer = self.__getattr__(f'layer{i + 1}')
            x = layer(x)
        return x

    def forward(self, x):
        return self._forward_impl(x)


class ResNetDBackbone(nn.Module, ABC):

    def __init__(self, in_channels=3, base_channels=64, layer_channels=(64, 128, 256, 512), layer_blocks=(2, 2, 2, 2), downsamples=(0, 1, 1, 1), groups=1, width_per_group=64, with_attentions=(0, 0, 0, 0), reduction=16, attention_type='SqueezeAndExcitationBlock2D', block_layer=None, conv_layer=None, norm_layer=None, act_layer=None, zero_init_residual=False, use_avg=False, fast_avg=False, **kwargs):
        """
        参考：[Bag of Tricks for Image Classification with Convolutional Neural Networks](https://arxiv.org/abs/1812.01187)
        :param in_channels: 输入通道数
        :param base_channels: 基础通道数
        :param layer_channels: 每一层通道数
        :param layer_blocks: 每一层块个数
        :param downsamples: 是否执行空间下采样
        :param groups: cardinality
        :param width_per_group: 每组的宽度
        :param with_attentions: 是否使用注意力模块
        :param reduction: 衰减率
        :param attention_type: 注意力模块类型
        :param block_layer: 块类型
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        :param zero_init_residual: 零初始化残差连接
        :param use_avg: 是否使用AvgPool进行下采样
        :param fast_avg: 在3x3之前执行下采样操作
        :param kwargs: 其他参数
        """
        super(ResNetDBackbone, self).__init__()
        assert len(layer_channels) == len(layer_blocks) == len(downsamples) == len(with_attentions)
        if block_layer is None:
            block_layer = BasicBlock
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        self.stem = make_stem(in_channels, base_channels, conv_layer, norm_layer, act_layer)
        in_channels = base_channels
        for i in range(len(layer_blocks)):
            res_layer = make_res_layer(in_channels, layer_channels[i], layer_blocks[i], downsamples[i], groups, width_per_group, with_attentions[i], reduction, attention_type, block_layer, conv_layer, norm_layer, act_layer, use_avg, fast_avg, **kwargs)
            in_channels = layer_channels[i] * block_layer.expansion
            layer_name = f'layer{i + 1}'
            self.add_module(layer_name, res_layer)
        self.layer_num = len(layer_blocks)
        self.init_weights(zero_init_residual)

    def init_weights(self, zero_init_residual):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck):
                    nn.init.constant_(m.bn3.weight, 0)
                elif isinstance(m, BasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)

    def _forward_impl(self, x):
        x = self.stem(x)
        for i in range(self.layer_num):
            layer = self.__getattr__(f'layer{i + 1}')
            x = layer(x)
        return x

    def forward(self, x):
        return self._forward_impl(x)


class SelectiveKernelConv2d(nn.Module, ABC):
    """
    refer to [Selective Kernel Networks](https://arxiv.org/abs/1903.06586)
    """

    def __init__(self, in_channels: int, out_channels: int, stride: int, groups: int, reduction_rate: int, default_channels: int=32, dimension: int=2):
        super().__init__()
        self.split1 = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, bias=False, padding=1, dilation=1, groups=groups), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True))
        self.split2 = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, bias=False, padding=2, dilation=2, groups=groups), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True))
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        inner_channels = max(out_channels // reduction_rate, default_channels)
        self.compact = nn.Sequential(nn.Conv2d(out_channels, inner_channels, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(inner_channels), nn.ReLU(inplace=True))
        self.select_a = nn.Linear(inner_channels, out_channels, bias=True)
        self.select_b = nn.Linear(inner_channels, out_channels, bias=True)
        self.softmax = nn.Softmax(dim=0)
        self.dimension = dimension
        self.out_channels = out_channels
        init_weights(self.modules())

    def forward(self, x: Tensor) ->Tensor:
        N, C = x.shape[:2]
        u1 = self.split1(x)
        u2 = self.split2(x)
        u = u1 + u2
        s = self.pool(u)
        z = self.compact(s).flatten(1)
        ac = self.select_a(z)
        bc = self.select_b(z)
        abc = torch.stack((ac, bc)).reshape(2, N, self.out_channels, *([1] * self.dimension))
        softmax_abc = self.softmax(abc)
        v = torch.sum(torch.stack((u1, u2)).mul(softmax_abc), dim=0)
        return v


class SKNetBlock(nn.Module, ABC):
    expansion = 4

    def __init__(self, in_channels, out_channels, stride=1, downsample=None, groups=1, base_width=64, with_attention=False, reduction=16, attention_type='SqueezeAndExcitationBlock2D', conv_layer=None, norm_layer=None, act_layer=None, **kwargs):
        """
        依次执行大小为1x1、3x3、1x1的卷积操作，如果进行下采样，那么使用第二个卷积层对输入空间尺寸进行减半操作
        参考Torchvision实现
        对于注意力模块，有两种嵌入方式：
        1. 对于Squeeze-And-Excitation或者Global Context操作，在残差连接中（after 1x1）嵌入；
        2. 对于NonLocal或者SimplifiedNonLoal，在Block完成计算后（after add）嵌入。
        对于Selective Kernel Conv2d，替换3x3卷积层
        :param in_channels: 输入通道数
        :param out_channels: 输出通道数
        :param stride: 步长
        :param downsample: 下采样
        :param groups: cardinality
        :param base_width: 基础宽度
        :param with_attention: 是否使用注意力模块
        :param reduction: 衰减率
        :param attention_type: 注意力模块类型
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        :param kwargs: 其他参数
        """
        super(SKNetBlock, self).__init__()
        assert with_attention in (0, 1)
        assert attention_type in ['GlobalContextBlock2D', 'SimplifiedNonLocal2DEmbeddedGaussian', 'NonLocal2DEmbeddedGaussian', 'SqueezeAndExcitationBlock2D']
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        self.down_sample = downsample
        width = int(out_channels * (base_width / 64.0)) * groups
        self.conv1 = conv_layer(in_channels, width, kernel_size=1, stride=1, bias=False)
        self.bn1 = norm_layer(width)
        self.conv2 = SelectiveKernelConv2d(width, width, stride=stride, groups=groups, reduction_rate=reduction)
        self.conv3 = conv_layer(width, out_channels * self.expansion, kernel_size=1, stride=1, bias=False)
        self.bn3 = norm_layer(out_channels * self.expansion)
        self.relu = act_layer(inplace=True)
        self.attention_after_1x1 = None
        self.attention_after_add = None
        if with_attention and attention_type in ['SqueezeAndExcitationBlock2D', 'GlobalContextBlock2D']:
            self.attention_after_1x1 = make_attention_block(out_channels * self.expansion, reduction, attention_type)
            self.attention_after_add = None
        if with_attention and attention_type in ['NonLocal2DEmbeddedGaussian', 'SimplifiedNonLocal2DEmbeddedGaussian']:
            self.attention_after_1x1 = None
            self.attention_after_add = make_attention_block(out_channels * self.expansion, reduction, attention_type)

    def forward(self, x):
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.relu(out)
        out = self.conv3(out)
        out = self.bn3(out)
        if self.attention_after_1x1 is not None:
            out = self.attention_after_1x1(out)
        if self.down_sample is not None:
            identity = self.down_sample(x)
        out += identity
        out = self.relu(out)
        if self.attention_after_add is not None:
            out = self.attention_after_add(out)
        return out


class ShuffleNetV1Unit(nn.Module, ABC):

    def __init__(self, in_channels, out_channels, groups, stride, downsample=None, with_group=True, conv_layer=None, norm_layer=None, act_layer=None):
        """
        In paper https://arxiv.org/abs/1707.01083
        when stride=1, Unit = ReLU(Add(
                                Input,
                                BN(1x1 GConv(
                                    BN(3x3 DWConv(
                                        Channel Shuffle(
                                            ReLU(BN(1x1 GConv(Input)))
                                        )
                                    ))
                                ))
                              ));
        when stride=2, Unit = ReLU(Concat(
                                S=2 K=3x3 AvgPool(Input),
                                BN(1x1 GConv(
                                  BN(S=2 K=3x3 DWConv(
                                    Channel Shuffle(
                                        ReLU(BN(1x1 GConv(Input)))
                                    )
                                  ))
                                ))
                              ));
        In official realization [ShuffleNet-Series/ShuffleNetV1/blocks.py](https://github.com/megvii-model/ShuffleNet-Series/blob/master/ShuffleNetV1/blocks.py)
        There are two differences with the paper description
        1. use channel shuffle after 3x3 DWConv;
        2. when stride=2, not use activation function for identity mapping
                                Unit = Concat(
                                    S=2 K=3x3 AvgPool(Input),
                                    ReLU(BN(1x1 GConv(
                                        Channel Shuffle(
                                            BN(S=2 K=3x3 DWConv(
                                                ReLU(BN(1x1 GConv(Input)))
                                            ))
                                        )
                                    )))
                                );
        for first:
        * [Why is ShuffleNetV1 different from description in paper ? #16](https://github.com/megvii-model/ShuffleNet-Series/issues/16)
        * [shufflenetv1的一个问题 #40](https://github.com/megvii-model/ShuffleNet-Series/issues/40)
        for second:
        * [A mismatch in shufflenetv1 about ReLU #53](https://github.com/megvii-model/ShuffleNet-Series/issues/53)
        :param in_channels: 输入通道
        :param out_channels: 输出通道
        :param groups: 分组数
        :param stride: 步长
        :param downsample: 作用于shortcut path
        :param with_group: 是否对第一个1x1逐点卷积应用分组
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        """
        super().__init__()
        assert out_channels % groups == 0
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        mid_channels = out_channels // 4
        out_channels = out_channels if stride == 1 else out_channels - in_channels
        self.conv1 = conv_layer(in_channels, mid_channels, kernel_size=1, stride=1, padding=0, bias=False, groups=groups if with_group else 1)
        self.norm1 = norm_layer(mid_channels)
        self.conv2 = conv_layer(mid_channels, mid_channels, kernel_size=3, stride=stride, padding=1, bias=False, groups=mid_channels)
        self.norm2 = norm_layer(mid_channels)
        self.conv3 = conv_layer(mid_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False, groups=groups)
        self.norm3 = norm_layer(out_channels)
        self.act = act_layer(inplace=True)
        self.down_sample = downsample
        self.groups = groups
        self.stride = stride

    def forward(self, x):
        identity = x
        out = self.conv1(x)
        out = self.norm1(out)
        out = self.act(out)
        out = self.conv2(out)
        out = self.norm2(out)
        if self.groups > 1:
            out = self.channel_shuffle(out)
        out = self.conv3(out)
        out = self.norm3(out)
        if self.stride == 2:
            out = torch.cat((self.down_sample(identity), out), dim=1)
        else:
            out = self.act(out + identity)
        return out

    def channel_shuffle(self, x):
        batchsize, num_channels, height, width = x.data.size()
        assert num_channels % self.groups == 0
        group_channels = num_channels // self.groups
        x = x.reshape(batchsize, group_channels, self.groups, height, width)
        x = x.permute(0, 2, 1, 3, 4)
        x = x.reshape(batchsize, num_channels, height, width)
        return x


def make_stage(in_channels, out_channels, block_num, with_downsample, block_layer, conv_layer, norm_layer, act_layer):
    """
    :param in_channels: 输入通道数
    :param out_channels: 输出通道数
    :param block_num: 块个数
    :param with_downsample: 是否执行空间下采样
    :param block_layer: 块类型
    :param conv_layer: 卷积层类型
    :param norm_layer: 归一化层类型
    :param act_layer: 激活层类型
    :return:
    """
    stride = 2 if with_downsample else 1
    if with_downsample:
        down_sample = nn.Sequential(conv_layer(in_channels, in_channels, kernel_size=3, stride=stride, padding=1, bias=False, groups=in_channels), norm_layer(in_channels), conv_layer(in_channels, out_channels // 2, kernel_size=1, stride=1, padding=0, bias=False), norm_layer(out_channels // 2), act_layer(inplace=True))
    else:
        down_sample = None
    blocks = list()
    blocks.append(block_layer(in_channels, out_channels, stride, down_sample, conv_layer, norm_layer, act_layer))
    in_channels = out_channels
    stride = 1
    down_sample = None
    for i in range(1, block_num):
        blocks.append(block_layer(in_channels // 2, out_channels, stride, down_sample, conv_layer, norm_layer, act_layer))
    return nn.Sequential(*blocks)


class ShuffleNetV1Backbone(nn.Module, ABC):

    def __init__(self, in_channels=3, base_channels=24, groups=8, stage_channels=(384, 768, 1536), stage_blocks=(4, 8, 4), downsamples=(1, 1, 1), with_groups=(0, 1, 1), block_layer=None, conv_layer=None, norm_layer=None, act_layer=None):
        """
        :param in_channels: 输入通道数
        :param base_channels: 第一个卷积层通道数
        :param groups: 分组数
        :param stage_channels: 每个阶段通道数
        :param stage_blocks: 每个阶段块个数
        :param downsamples: 是否执行空间下采样
        :param with_groups: 是否对第一个1x1逐点卷积执行分组操作
        :param block_layer: 块类型
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        """
        super(ShuffleNetV1Backbone, self).__init__()
        if block_layer is None:
            block_layer = ShuffleNetV1Unit
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        self.stem = make_stem(in_channels, base_channels, conv_layer, norm_layer, act_layer)
        in_channels = base_channels
        for i in range(len(stage_blocks)):
            res_layer = make_stage(in_channels, stage_channels[i], groups, stage_blocks[i], downsamples[i], with_groups[i], block_layer, conv_layer, norm_layer, act_layer)
            in_channels = stage_channels[i]
            stage_name = f'stage{i + 1}'
            self.add_module(stage_name, res_layer)
        self.stage_num = len(stage_blocks)
        init_weights(self.modules())

    def forward(self, x):
        x = self.stem(x)
        for i in range(self.stage_num):
            stage = self.__getattr__(f'stage{i + 1}')
            x = stage(x)
        return x


def channel_shuffle(x, groups):
    """
    # >>> a = torch.arange(12)
    # >>> b = a.reshape(3,4)
    # >>> c = b.transpose(1,0).contiguous()
    # >>> d = c.view(3,4)
    # >>> a
    # tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11])
    # >>> b
    # tensor([[ 0,  1,  2,  3],
    #         [ 4,  5,  6,  7],
    #         [ 8,  9, 10, 11]])
    # >>> c
    # tensor([[ 0,  4,  8],
    #         [ 1,  5,  9],
    #         [ 2,  6, 10],
    #         [ 3,  7, 11]])
    # >>> d
    # tensor([[ 0,  4,  8,  1],
    #         [ 5,  9,  2,  6],
    #         [10,  3,  7, 11]])
    """
    batchsize, num_channels, height, width = x.data.size()
    channels_per_group = num_channels // groups
    x = x.view(batchsize, groups, channels_per_group, height, width)
    x = torch.transpose(x, 1, 2).contiguous()
    x = x.view(batchsize, -1, height, width)
    return x


class ShuffleNetV2Unit(nn.Module, ABC):

    def __init__(self, in_channels, out_channels, stride, downsample=None, conv_layer=None, norm_layer=None, act_layer=None):
        """
        refer to Torchvision realization: https://github.com/pytorch/vision/blob/master/torchvision/models/shufflenetv2.py
        when stride = 1, Unit = Channel Shuffle(Concat(Channel Split(Input), Conv(DWConv(Conv(Channel Split(Input))))));
        when stride = 2, Unit = Channel Shuffle(Concat(Conv(DWConv(Input)), Conv(DWConv(Conv(Input)))));
        official realization (https://github.com/megvii-model/ShuffleNet-Series/blob/master/ShuffleNetV2/blocks.py) has
         a more ingenious implementation, see [A mismatch in shufflenetv1 about ReLU #53](https://github.com/megvii-model/ShuffleNet-Series/issues/53)
        :param in_channels: 输入通道
        :param out_channels: 输出通道
        :param stride: 步长
        :param downsample: 作用于shortcut path
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        """
        super().__init__()
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        mid_channels = out_channels // 2
        self.conv1 = conv_layer(in_channels, mid_channels, kernel_size=1, stride=1, padding=0, bias=False)
        self.norm1 = norm_layer(mid_channels)
        self.conv2 = conv_layer(mid_channels, mid_channels, kernel_size=3, stride=stride, padding=1, bias=False, groups=mid_channels)
        self.norm2 = norm_layer(mid_channels)
        self.conv3 = conv_layer(mid_channels, mid_channels, kernel_size=1, stride=1, padding=0, bias=False)
        self.norm3 = norm_layer(mid_channels)
        self.act = act_layer(inplace=True)
        self.stride = stride
        self.down_sample = downsample

    def forward(self, x):
        if self.stride == 1:
            x1, x2 = x.chunk(2, dim=1)
        else:
            x1 = x
            x2 = x
        out = self.conv1(x2)
        out = self.norm1(out)
        out = self.act(out)
        out = self.conv2(out)
        out = self.norm2(out)
        out = self.conv3(out)
        out = self.norm3(out)
        out = self.act(out)
        if self.down_sample is not None:
            x1 = self.down_sample(x1)
        out = torch.cat((x1, out), dim=1)
        out = channel_shuffle(out, 2)
        return out


class ShuffleNetV2Backbone(nn.Module, ABC):

    def __init__(self, in_channels=3, base_channels=24, out_channels=1024, stage_channels=(116, 232, 464), stage_blocks=(4, 8, 4), downsamples=(1, 1, 1), block_layer=None, conv_layer=None, norm_layer=None, act_layer=None):
        """
        :param in_channels: 输入通道数
        :param base_channels: 第一个卷积层通道数
        :param out_channels: 输出通道数
        :param stage_channels: 每一层通道数
        :param stage_blocks: 每一层块个数
        :param downsamples: 是否执行空间下采样
        :param block_layer: 块类型
        :param conv_layer: 卷积层类型
        :param norm_layer: 归一化层类型
        :param act_layer: 激活层类型
        """
        super(ShuffleNetV2Backbone, self).__init__()
        if block_layer is None:
            block_layer = ShuffleNetV2Unit
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if act_layer is None:
            act_layer = nn.ReLU
        self.first_stem, self.last_stem = make_stem(in_channels, base_channels, stage_channels[-1], out_channels, conv_layer, norm_layer, act_layer)
        in_channels = base_channels
        for i in range(len(stage_blocks)):
            res_layer = make_stage(in_channels, stage_channels[i], stage_blocks[i], downsamples[i], block_layer, conv_layer, norm_layer, act_layer)
            in_channels = stage_channels[i]
            layer_name = f'stage{i + 1}'
            self.add_module(layer_name, res_layer)
        self.stage_num = len(stage_blocks)
        init_weights(self.modules())

    def forward(self, x):
        x = self.first_stem(x)
        for i in range(self.stage_num):
            stage = self.__getattr__(f'stage{i + 1}')
            x = stage(x)
        x = self.last_stem(x)
        return x


class RepVGGBackbone(nn.Module):

    def __init__(self, in_channels=3, base_channels=64, layer_channels=(64, 128, 256, 512), layer_blocks=(2, 4, 14, 1), width_multipliers=(1.0, 1.0), down_samples=(1, 1, 1, 1), groups=dict(), with_attention=False, reduction=16, attention_type='SqueezeAndExcitationBlock2D', attention_bias=False, conv_layer=None, act_layer=None):
        """
        :param in_channels: 输入通道数
        :param base_channels: 基础通道数
        :param layer_channels: 每一层通道数
        :param layer_blocks: 每一层块个数
        :param width_multipliers: 宽度乘法器
        :param down_samples: 是否执行空间下采样
        :param groups: cardinality
        :param conv_layer: 卷积层类型
        :param act_layer: 激活层类型
        """
        super(RepVGGBackbone, self).__init__()
        assert len(layer_channels) == len(layer_blocks) == len(down_samples)
        assert len(width_multipliers) == 2
        assert isinstance(groups, dict)
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if act_layer is None:
            act_layer = nn.ReLU
        self.with_attention = with_attention
        self.reduction = reduction
        self.attention_type = attention_type
        self.attention_bias = attention_bias
        self.cur_layer_idx = 0
        self.override_groups_map = groups
        cur_groups = self.override_groups_map.get(self.cur_layer_idx, 1)
        width_multiplier_a, width_multiplier_b = width_multipliers
        base_channels = min(int(base_channels), int(base_channels * width_multiplier_a))
        self.stage0 = make_stem(in_channels, base_channels, cur_groups, with_attention, reduction, attention_type, attention_bias, conv_layer, act_layer)
        in_channels = base_channels
        for i in range(len(layer_blocks)):
            width_multiplier = width_multiplier_a if i != len(layer_blocks) - 1 else width_multiplier_b
            out_channels = int(layer_channels[i] * width_multiplier)
            stage = self._make_stage(in_channels, out_channels, layer_blocks[i], down_samples[i], conv_layer, act_layer)
            in_channels = out_channels
            stage_name = f'stage{i + 1}'
            self.add_module(stage_name, stage)
        init_weights(self.modules())

    def _make_stage(self, in_channels, out_channels, block_num, with_sample, conv_layer, act_layer):
        """
        :param in_channels: 输入通道数
        :param out_channels: 输出通道数
        :param block_num: 块个数
        :param with_sample: 是否执行空间下采样
        :param conv_layer: 卷积层类型
        :param act_layer: 激活层类型
        :return:
        """
        stride = 2 if with_sample else 1
        padding = 1 if stride == 2 else 0
        self.cur_layer_idx += 1
        cur_groups = self.override_groups_map.get(self.cur_layer_idx, 1)
        blocks = list()
        blocks.append(nn.Sequential(conv_layer(in_channels, out_channels, kernel_size=3, stride=stride, padding=padding, groups=cur_groups, bias=True), make_attention_block(out_channels, self.reduction, self.attention_type, bias=self.attention_bias) if self.with_attention else nn.Identity(), act_layer(inplace=True)))
        in_channels = out_channels
        stride = 1
        padding = 1
        for i in range(1, block_num):
            self.cur_layer_idx += 1
            cur_groups = self.override_groups_map.get(self.cur_layer_idx, 1)
            blocks.append(nn.Sequential(conv_layer(in_channels, out_channels, kernel_size=3, stride=stride, padding=padding, groups=cur_groups, bias=True), make_attention_block(out_channels, self.reduction, self.attention_type, bias=self.attention_bias) if self.with_attention else nn.Identity(), act_layer(inplace=True)))
        return nn.Sequential(*blocks)

    def _forward_impl(self, x):
        x = self.stage0(x)
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)
        x = self.stage4(x)
        return x

    def forward(self, x):
        return self._forward_impl(x)


KEY_LOSS = 'loss'


KEY_OUTPUT = 'output'


class CrossEntropyLoss(nn.Module, ABC):

    def __init__(self, cfg):
        super(CrossEntropyLoss, self).__init__()
        self.loss = nn.CrossEntropyLoss(reduction='mean')

    def __call__(self, output_dict, targets):
        assert isinstance(output_dict, dict) and KEY_OUTPUT in output_dict.keys()
        inputs = output_dict[KEY_OUTPUT]
        loss = self.loss(inputs, targets)
        return {KEY_LOSS: loss}


class LabelSmoothingLoss(nn.Module, ABC):
    """
    Label smoothing cross entropy loss
    Refer to:
    1. [解决过拟合：如何在PyTorch中使用标签平滑正则化](https://zhuanlan.zhihu.com/p/123077402)
    2. [Rethinking the Inception Architecture for Computer Vision](https://arxiv.org/abs/1512.00567)
    3. [[译]Rethinking the Inception Architecture for Computer Vision](https://blog.zhujian.life/posts/a0a2be91.html)
    4. [Label Smoothing in PyTorch](https://stackoverflow.com/questions/55681502/label-smoothing-in-pytorch)

    from 1.10.0, torch add label_smoothing_loss to CrossEntropyLoss
    """

    def __init__(self, epsilon, reduction='mean'):
        super(LabelSmoothingLoss, self).__init__()
        self.epsilon = epsilon
        self.reduction = reduction

    def __call__(self, output_dict, targets):
        assert isinstance(output_dict, dict) and KEY_OUTPUT in output_dict.keys()
        inputs = output_dict[KEY_OUTPUT]
        n = inputs.size()[-1]
        log_preds = F.log_softmax(inputs, dim=-1)
        loss = self.reduce_loss(-log_preds.sum(dim=-1), reduction=self.reduction)
        nll = F.nll_loss(log_preds, targets, reduction=self.reduction)
        return {KEY_LOSS: self.linear_combination(loss / n, nll, self.epsilon)}

    def reduce_loss(self, loss, reduction='mean'):
        return loss.mean() if reduction == 'mean' else loss.sum() if reduction == 'sum' else loss

    def linear_combination(self, x, y, epsilon):
        return epsilon * x + (1 - epsilon) * y


class GeneralHead2D(nn.Module, ABC):

    def __init__(self, feature_dims=1024, dropout_rate=0.0, num_classes=1000, bias=True):
        """
        AvgPool + Dropout + FC
        :param feature_dims: 输入特征维度
        :param dropout_rate: 随机失活概率
        :param num_classes: 类别数
        """
        super(GeneralHead2D, self).__init__()
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.dropout = nn.Dropout(p=dropout_rate)
        self.fc = nn.Linear(feature_dims, num_classes, bias=bias)
        self.init_weights()

    def init_weights(self):
        nn.init.normal_(self.fc.weight, 0, 0.01)
        if self.fc.bias is not None:
            nn.init.zeros_(self.fc.bias)

    def forward(self, x):
        x = self.pool(x)
        x = torch.flatten(x, 1)
        x = self.dropout(x)
        x = self.fc(x)
        return x


class GeneralHead3D(nn.Module, ABC):

    def __init__(self, feature_dims=2048, dropout_rate=0.0, num_classes=1000):
        """
        AvgPool + Dropout + FC
        :param feature_dims: 输入特征维度
        :param dropout_rate: 随机失活概率
        :param num_classes: 类别数
        """
        super(GeneralHead3D, self).__init__()
        self.pool = nn.AdaptiveAvgPool3d((1, 1, 1))
        self.dropout = nn.Dropout(p=dropout_rate)
        self.fc = nn.Linear(feature_dims, num_classes)
        self.init_weights()

    def init_weights(self):
        nn.init.normal_(self.fc.weight, 0, 0.01)
        nn.init.zeros_(self.fc.bias)

    def forward(self, x):
        x = self.pool(x)
        x = torch.flatten(x, 1)
        x = self.dropout(x)
        x = self.fc(x)
        return x


class MobileNetV3Head(nn.Module, ABC):

    def __init__(self, feature_dims=960, inner_dims=1280, dropout_rate=0.0, num_classes=1000, conv_layer=None, act_layer=None):
        """
        :param feature_dims: 输入特征维度
        :param inner_dims: 中间特征维度
        :param dropout_rate: 随机失活概率
        :param num_classes: 类别数
        :param conv_layer: 卷积层类型
        :param act_layer: 激活层类型
        """
        super(MobileNetV3Head, self).__init__()
        if conv_layer is None:
            conv_layer = nn.Conv2d
        if act_layer is None:
            act_layer = nn.Hardswish
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.conv1 = conv_layer(feature_dims, inner_dims, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), bias=True)
        self.conv2 = conv_layer(inner_dims, num_classes, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), bias=True)
        self.act = act_layer(inplace=True)
        self.dropout = nn.Dropout(p=dropout_rate)
        self.init_weights()

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def forward(self, x):
        x = self.pool(x)
        x = self.conv1(x)
        x = self.act(x)
        x = self.dropout(x)
        x = self.conv2(x)
        x = torch.flatten(x, 1)
        return x


class CropLayer(nn.Module):

    def __init__(self, crop_set):
        super(CropLayer, self).__init__()
        self.rows_to_crop = -crop_set[0]
        self.cols_to_crop = -crop_set[1]
        assert self.rows_to_crop >= 0
        assert self.cols_to_crop >= 0

    def forward(self, input):
        if self.rows_to_crop == 0 and self.cols_to_crop == 0:
            return input
        elif self.rows_to_crop > 0 and self.cols_to_crop == 0:
            return input[:, :, self.rows_to_crop:-self.rows_to_crop, :]
        elif self.rows_to_crop == 0 and self.cols_to_crop > 0:
            return input[:, :, :, self.cols_to_crop:-self.cols_to_crop]
        else:
            return input[:, :, self.rows_to_crop:-self.rows_to_crop, self.cols_to_crop:-self.cols_to_crop]


class AsymmetricConvolutionBlock(nn.Module):
    """
    refer to [ACNet/acnet/acb.py](https://github.com/DingXiaoH/ACNet/blob/master/acnet/acb.py)
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, padding_mode='zeros', use_affine=True, use_last_bn=False):
        super(AsymmetricConvolutionBlock, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.groups = groups
        self.padding_mode = padding_mode
        self.square_conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(kernel_size, kernel_size), stride=stride, padding=padding, dilation=dilation, groups=groups, bias=False, padding_mode=padding_mode)
        self.square_bn = nn.BatchNorm2d(num_features=out_channels, affine=use_affine)
        center_offset_from_origin_border = padding - kernel_size // 2
        ver_pad_or_crop = padding, center_offset_from_origin_border
        hor_pad_or_crop = center_offset_from_origin_border, padding
        if center_offset_from_origin_border >= 0:
            self.ver_conv_crop_layer = nn.Identity()
            ver_conv_padding = ver_pad_or_crop
            self.hor_conv_crop_layer = nn.Identity()
            hor_conv_padding = hor_pad_or_crop
        else:
            self.ver_conv_crop_layer = CropLayer(crop_set=ver_pad_or_crop)
            ver_conv_padding = 0, 0
            self.hor_conv_crop_layer = CropLayer(crop_set=hor_pad_or_crop)
            hor_conv_padding = 0, 0
        self.ver_conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(kernel_size, 1), stride=stride, padding=ver_conv_padding, dilation=dilation, groups=groups, bias=False, padding_mode=padding_mode)
        self.hor_conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(1, kernel_size), stride=stride, padding=hor_conv_padding, dilation=dilation, groups=groups, bias=False, padding_mode=padding_mode)
        self.ver_bn = nn.BatchNorm2d(num_features=out_channels, affine=use_affine)
        self.hor_bn = nn.BatchNorm2d(num_features=out_channels, affine=use_affine)
        if use_last_bn:
            self.last_bn = nn.BatchNorm2d(num_features=out_channels, affine=True)
        self._init_weights()

    def _init_weights(self, gamma=1.0 / 3):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, gamma)
                nn.init.constant_(m.bias, gamma)

    def forward(self, input):
        square_outputs = self.square_conv(input)
        square_outputs = self.square_bn(square_outputs)
        vertical_outputs = self.ver_conv_crop_layer(input)
        vertical_outputs = self.ver_conv(vertical_outputs)
        vertical_outputs = self.ver_bn(vertical_outputs)
        horizontal_outputs = self.hor_conv_crop_layer(input)
        horizontal_outputs = self.hor_conv(horizontal_outputs)
        horizontal_outputs = self.hor_bn(horizontal_outputs)
        result = square_outputs + vertical_outputs + horizontal_outputs
        if hasattr(self, 'last_bn'):
            return self.last_bn(result)
        return result


class BatchNorm2d(nn.BatchNorm2d):
    """
    from [pytorch_misc/batch_norm_manual.py](https://github.com/ptrblck/pytorch_misc/blob/master/batch_norm_manual.py)
    """

    def __init__(self, num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True):
        super(BatchNorm2d, self).__init__(num_features, eps, momentum, affine, track_running_stats)

    def forward(self, input):
        self._check_input_dim(input)
        exponential_average_factor = 0.0
        if self.training and self.track_running_stats:
            if self.num_batches_tracked is not None:
                self.num_batches_tracked += 1
                if self.momentum is None:
                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)
                else:
                    exponential_average_factor = self.momentum
        if self.training:
            mean = input.mean([0, 2, 3])
            var = input.var([0, 2, 3], unbiased=False)
            n = input.numel() / input.size(1)
            with torch.no_grad():
                self.running_mean = exponential_average_factor * mean + (1 - exponential_average_factor) * self.running_mean
                self.running_var = exponential_average_factor * var * n / (n - 1) + (1 - exponential_average_factor) * self.running_var
        else:
            mean = self.running_mean
            var = self.running_var
        input = (input - mean[None, :, None, None]) / torch.sqrt(var[None, :, None, None] + self.eps)
        if self.affine:
            input = input * self.weight[None, :, None, None] + self.bias[None, :, None, None]
        return input


class IdentityBasedConv1x1(nn.Conv2d):

    def __init__(self, channels, groups=1):
        super(IdentityBasedConv1x1, self).__init__(in_channels=channels, out_channels=channels, kernel_size=(1, 1), stride=(1, 1), padding=0, groups=groups, bias=False)
        assert channels % groups == 0
        input_dim = channels // groups
        id_value = np.zeros((channels, input_dim, 1, 1))
        for i in range(channels):
            id_value[i, i % input_dim, 0, 0] = 1
        self.id_tensor = torch.from_numpy(id_value).type_as(self.weight)
        nn.init.zeros_(self.weight)

    def forward(self, input):
        kernel = self.weight + self.id_tensor
        result = F.conv2d(input, kernel, None, stride=1, padding=0, dilation=self.dilation, groups=self.groups)
        return result

    def get_actual_kernel(self):
        return self.weight + self.id_tensor


class BNAndPadLayer(nn.Module):

    def __init__(self, pad_pixels, num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True):
        super(BNAndPadLayer, self).__init__()
        self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine, track_running_stats)
        self.pad_pixels = pad_pixels

    def forward(self, input):
        output = self.bn(input)
        if self.pad_pixels > 0:
            if self.bn.affine:
                pad_values = self.bn.bias.detach() - self.bn.running_mean * self.bn.weight.detach() / torch.sqrt(self.bn.running_var + self.bn.eps)
            else:
                pad_values = -self.bn.running_mean / torch.sqrt(self.bn.running_var + self.bn.eps)
            output = F.pad(output, [self.pad_pixels] * 4)
            pad_values = pad_values.view(1, -1, 1, 1)
            output[:, :, 0:self.pad_pixels, :] = pad_values
            output[:, :, -self.pad_pixels:, :] = pad_values
            output[:, :, :, 0:self.pad_pixels] = pad_values
            output[:, :, :, -self.pad_pixels:] = pad_values
        return output

    @property
    def weight(self):
        return self.bn.weight

    @property
    def bias(self):
        return self.bn.bias

    @property
    def running_mean(self):
        return self.bn.running_mean

    @property
    def running_var(self):
        return self.bn.running_var

    @property
    def eps(self):
        return self.bn.eps


def conv_bn(in_channels, out_channels, kernel_size, stride, padding, groups=1):
    result = nn.Sequential()
    result.add_module('conv', nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, bias=False))
    result.add_module('bn', nn.BatchNorm2d(num_features=out_channels))
    return result


class DiverseBranchBlock(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, internal_channels_1x1_3x3=None, single_init=False):
        super(DiverseBranchBlock, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.groups = groups
        assert padding == kernel_size // 2
        self.dbb_origin = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups)
        self.dbb_avg = nn.Sequential()
        if groups < out_channels:
            self.dbb_avg.add_module('conv', nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(1, 1), stride=(1, 1), padding=0, groups=groups, bias=False))
            self.dbb_avg.add_module('bn', BNAndPadLayer(pad_pixels=padding, num_features=out_channels))
            self.dbb_avg.add_module('avg', nn.AvgPool2d(kernel_size=kernel_size, stride=stride, padding=0))
            self.dbb_1x1 = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=0, groups=groups)
        else:
            self.dbb_avg.add_module('avg', nn.AvgPool2d(kernel_size=kernel_size, stride=stride, padding=padding))
        self.dbb_avg.add_module('avgbn', nn.BatchNorm2d(out_channels))
        if internal_channels_1x1_3x3 is None:
            internal_channels_1x1_3x3 = in_channels if groups < out_channels else 2 * in_channels
        self.dbb_1x1_kxk = nn.Sequential()
        if internal_channels_1x1_3x3 == in_channels:
            self.dbb_1x1_kxk.add_module('idconv1', IdentityBasedConv1x1(channels=in_channels, groups=groups))
        else:
            self.dbb_1x1_kxk.add_module('conv1', nn.Conv2d(in_channels=in_channels, out_channels=internal_channels_1x1_3x3, kernel_size=(1, 1), stride=(1, 1), padding=0, groups=groups, bias=False))
        self.dbb_1x1_kxk.add_module('bn1', BNAndPadLayer(pad_pixels=padding, num_features=internal_channels_1x1_3x3, affine=True))
        self.dbb_1x1_kxk.add_module('conv2', nn.Conv2d(in_channels=internal_channels_1x1_3x3, out_channels=out_channels, kernel_size=kernel_size, stride=(stride, stride), padding=0, groups=groups, bias=False))
        self.dbb_1x1_kxk.add_module('bn2', nn.BatchNorm2d(out_channels))
        self.init_weights()
        if single_init:
            self.single_init()

    def forward(self, inputs):
        out = self.dbb_origin(inputs)
        if hasattr(self, 'dbb_1x1'):
            out += self.dbb_1x1(inputs)
        out += self.dbb_avg(inputs)
        out += self.dbb_1x1_kxk(inputs)
        return out

    def init_weights(self):
        if hasattr(self, 'dbb_origin'):
            init_helper.init_weights(self.dbb_origin)
        if hasattr(self, 'dbb_1x1'):
            init_helper.init_weights(self.dbb_1x1)
        if hasattr(self, 'dbb_avg'):
            init_helper.init_weights(self.dbb_avg)
        if hasattr(self, 'dbb_1x1_kxk'):
            init_helper.init_weights(self.dbb_1x1_kxk)

    def init_gamma(self, gamma_value):
        if hasattr(self, 'dbb_origin'):
            torch.nn.init.constant_(self.dbb_origin.bn.weight, gamma_value)
        if hasattr(self, 'dbb_1x1'):
            torch.nn.init.constant_(self.dbb_1x1.bn.weight, gamma_value)
        if hasattr(self, 'dbb_avg'):
            torch.nn.init.constant_(self.dbb_avg.avgbn.weight, gamma_value)
        if hasattr(self, 'dbb_1x1_kxk'):
            torch.nn.init.constant_(self.dbb_1x1_kxk.bn2.weight, gamma_value)

    def single_init(self):
        self.init_gamma(0.0)
        if hasattr(self, 'dbb_origin'):
            torch.nn.init.constant_(self.dbb_origin.bn.weight, 1.0)


class GlobalContextBlock1D(_GlobalContextBlockND):

    def __init__(self, in_channels, reduction=16, dimension=1):
        super().__init__(in_channels, reduction, dimension)

    def _check_input_dim(self, input):
        if input.dim() != 2 and input.dim() != 3:
            raise ValueError('expected 2D or 3D input (got {}D input)'.format(input.dim()))


class GlobalContextBlock3D(_GlobalContextBlockND):

    def __init__(self, in_channels, reduction=16, dimension=3):
        super().__init__(in_channels, reduction, dimension)

    def _check_input_dim(self, input):
        if input.dim() != 5:
            raise ValueError('expected 5D input (got {}D input)'.format(input.dim()))


class GroupNorm(nn.Module):
    """
    refer to [group_normalization/group_norm.py](https://github.com/taokong/group_normalization/blob/master/group_norm.py)
    """

    def __init__(self, num_groups: int=4, num_channels: int=32, eps: float=1e-05, affine: bool=True):
        super(GroupNorm, self).__init__()
        self.num_groups = num_groups
        self.num_channels = num_channels
        self.eps = eps
        self.affine = affine
        if affine:
            self.weight = torch.ones(num_channels, 1, 1)
            self.bias = torch.zeros(num_channels, 1, 1)

    def forward(self, inputs):
        shape = inputs.size()
        N, C = shape[:2]
        x = inputs.view(N, self.num_groups, -1)
        mean = x.mean(dim=2, keepdim=True)
        std = x.std(dim=2, keepdim=True)
        x = (x - mean) / (std + self.eps)
        x = x.view(shape)
        return x * self.weight + self.bias


class GroupNormWrapper(nn.GroupNorm):

    def __init__(self, num_channels: int, num_groups: int=4, eps: float=1e-05, affine: bool=True) ->None:
        super(GroupNormWrapper, self).__init__(num_groups, num_channels, eps, affine)
        self.num_features = num_channels


class LayerNorm(nn.LayerNorm):

    def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=False):
        super(LayerNorm, self).__init__(normalized_shape, eps, elementwise_affine)

    def forward(self, input):
        mean = input.mean([1, 2, 3])
        var = input.var([1, 2, 3], unbiased=False)
        input = (input - mean[:, None, None, None]) / torch.sqrt(var[:, None, None, None] + self.eps)
        if self.elementwise_affine:
            input = input * self.weight[None, :] + self.bias[None, :]
        return input


class NonLocal1DEmbeddedGaussian(_NonLocalNDEmbeddedGaussian):

    def __init__(self, in_channels, inner_channels=None, dimension=1, with_pool=True, norm_layer=None, zero_init_final_norm=True):
        super().__init__(in_channels, inner_channels, dimension, with_pool, norm_layer, zero_init_final_norm)

    def _check_input_dim(self, input):
        if input.dim() != 2 and input.dim() != 3:
            raise ValueError('expected 2D or 3D input (got {}D input)'.format(input.dim()))


class NonLocal3DEmbeddedGaussian(_NonLocalNDEmbeddedGaussian):

    def __init__(self, in_channels, inner_channels=None, dimension=3, with_pool=True, norm_layer=None, zero_init_final_norm=True):
        super().__init__(in_channels, inner_channels, dimension, with_pool, norm_layer, zero_init_final_norm)

    def _check_input_dim(self, input):
        if input.dim() != 5:
            raise ValueError('expected 5D input (got {}D input)'.format(input.dim()))


class RepVGGBlock(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, padding_mode='zeros'):
        super(RepVGGBlock, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.groups = groups
        self.padding_mode = padding_mode
        padding_11 = padding - kernel_size // 2
        self.rbr_identity = nn.BatchNorm2d(num_features=in_channels) if out_channels == in_channels and stride == 1 else None
        self.rbr_dense = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups)
        self.rbr_1x1 = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=padding_11, groups=groups)
        self._init_weights()

    def _init_weights(self, gamma=0.01):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, gamma)
                nn.init.constant_(m.bias, gamma)

    def forward(self, inputs):
        if self.rbr_identity is None:
            id_out = 0
        else:
            id_out = self.rbr_identity(inputs)
        return self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out

    def repvgg_convert(self):
        kernel, bias = self.get_equivalent_kernel_bias()
        return kernel.detach().cpu().numpy(), bias.detach().cpu().numpy()


class SimplifiedNonLocal1DEmbeddedGaussian(_SimplifiedNonLocalNDEmbeddedGaussian):

    def __init__(self, in_channels, dimension=1, norm_layer=None, zero_init_final_norm=True):
        super().__init__(in_channels, dimension, norm_layer, zero_init_final_norm)

    def _check_input_dim(self, input):
        if input.dim() != 2 and input.dim() != 3:
            raise ValueError('expected 2D or 3D input (got {}D input)'.format(input.dim()))


class SimplifiedNonLocal3DEmbeddedGaussian(_SimplifiedNonLocalNDEmbeddedGaussian):

    def __init__(self, in_channels, dimension=3, norm_layer=None, zero_init_final_norm=True):
        super().__init__(in_channels, dimension, norm_layer, zero_init_final_norm)

    def _check_input_dim(self, input):
        if input.dim() != 5:
            raise ValueError('expected 5D input (got {}D input)'.format(input.dim()))


class SqueezeAndExcitationBlock1D(_SqueezeAndExcitationBlockND):

    def __init__(self, in_channels, reduction=16, dimension=1, **kwargs):
        super().__init__(in_channels, reduction, dimension, **kwargs)

    def _check_input_dim(self, input):
        if input.dim() != 2 and input.dim() != 3:
            raise ValueError('expected 2D or 3D input (got {}D input)'.format(input.dim()))


class SqueezeAndExcitationBlock3D(_SqueezeAndExcitationBlockND):

    def __init__(self, in_channels, reduction=16, dimension=3, **kwargs):
        super().__init__(in_channels, reduction, dimension, **kwargs)

    def _check_input_dim(self, input):
        if input.dim() != 5:
            raise ValueError('expected 5D input (got {}D input)'.format(input.dim()))


def build_backbone(cfg):
    return registry.Backbone[cfg.MODEL.BACKBONE.NAME](cfg)


def build_head(cfg):
    return registry.HEAD[cfg.MODEL.HEAD.NAME](cfg)


def freezing_bn(model, partial_bn=False):
    count = 0
    for m in model.modules():
        if isinstance(m, nn.BatchNorm2d):
            count += 1
            if count == 1 and partial_bn:
                continue
            m.eval()
            m.weight.requires_grad = False
            m.bias.requires_grad = False


def load_pretrained_weights(model, model_name, weights_path=None, load_fc=True, verbose=True, url_map=None):
    """Loads pretrained weights from weights path or download using url.

    Args:
        model (Module): The whole model.
        model_name (str): Model name.
        weights_path (None or str):
            str: path to pretrained weights file on the local disk.
            None: use pretrained weights downloaded from the Internet.
        load_fc (bool): Whether to load pretrained weights for fc layer at the end of the model.
        url_map (dict or None): Remote pre-training model path corresponding to each model name
    """
    assert url_map is None or isinstance(url_map, dict), url_map
    assert weights_path is None or os.path.isfile(weights_path), weights_path
    if isinstance(weights_path, str):
        state_dict = torch.load(weights_path)['model']
    elif url_map is not None:
        if model_name not in url_map.keys() or url_map[model_name] == '':
            if verbose:
                None
            return None
        remote_url = url_map[model_name]
        state_dict = model_zoo.load_url(remote_url)['model']
    else:
        return None
    if load_fc:
        ret = model.load_state_dict(state_dict, strict=False)
        assert not ret.missing_keys, 'Missing keys when loading pretrained weights: {}'.format(ret.missing_keys)
    else:
        missing_keys_list = list()
        if 'head.fc.weight' in state_dict.keys():
            state_dict.pop('head.fc.weight')
            missing_keys_list.append('head.fc.weight')
        if 'head.fc.bias' in state_dict.keys():
            state_dict.pop('head.fc.bias')
            missing_keys_list.append('head.fc.bias')
        if 'head.conv2.weight' in state_dict.keys():
            state_dict.pop('head.conv2.weight')
        if 'head.conv2.bias' in state_dict.keys():
            state_dict.pop('head.conv2.bias')
        ret = model.load_state_dict(state_dict, strict=False)
        assert set(ret.missing_keys) == set(missing_keys_list), 'Missing keys when loading pretrained weights: {}'.format(ret.missing_keys)
    if verbose:
        None


class BaseRecognizer(nn.Module, ABC):

    def __init__(self, cfg):
        super(BaseRecognizer, self).__init__()
        self.fix_bn = cfg.MODEL.NORM.FIX_BN
        self.partial_bn = cfg.MODEL.NORM.PARTIAL_BN
        self.backbone = build_backbone(cfg)
        self.head = build_head(cfg)
        self._init_weights(cfg)

    def _init_weights(self, cfg):
        pretrained_local = cfg.MODEL.RECOGNIZER.PRETRAINED_LOCAL
        pretrained_num_classes = cfg.MODEL.RECOGNIZER.PRETRAINED_NUM_CLASSES
        num_classes = cfg.MODEL.HEAD.NUM_CLASSES
        load_pretrained_weights(self, cfg.MODEL.BACKBONE.ARCH, weights_path=None if pretrained_local == '' else pretrained_local, load_fc=pretrained_num_classes == num_classes, verbose=True, url_map=None)

    def train(self, mode: bool=True) ->T:
        super(BaseRecognizer, self).train(mode=mode)
        if mode and (self.partial_bn or self.fix_bn):
            freezing_bn(self, partial_bn=self.partial_bn)
        return self

    def forward(self, x):
        x = self.backbone(x)
        x = self.head(x)
        return {KEY_OUTPUT: x}


class MobileNetV3(BaseRecognizer):

    def __init__(self, cfg):
        super().__init__(cfg)

    def init_weights(self, pretrained, pretrained_num_classes, num_classes):
        super(MobileNetV3, self).init_weights(pretrained, pretrained_num_classes, pretrained_num_classes)
        if num_classes != pretrained_num_classes:
            in_channels = self.head.conv2.in_channels
            conv2 = nn.Conv2d(in_channels, num_classes, kernel_size=1, stride=1, padding=0, bias=True)
            nn.init.kaiming_normal_(conv2.weight, mode='fan_out', nonlinearity='relu')
            nn.init.zeros_(conv2.bias)
            self.head.conv2 = conv2


class TorchvisionMNASNet(nn.Module, ABC):

    def __init__(self, width_multiplier=1.0, num_classes=1000, pretrained=False, pretrained_num_classes=1000, fix_bn=False, partial_bn=False):
        """
        :param width_multiplier: 宽度乘法器
        :param num_classes: 类别数
        :param pretrained: 预训练模型
        :param pretrained_num_classes: 假定预训练模型类别数
        :param fix_bn: 固定BN
        :param partial_bn: 仅训练第一层BN
        """
        super(TorchvisionMNASNet, self).__init__()
        self.fix_bn = fix_bn
        self.partial_bn = partial_bn
        if width_multiplier == 0.5:
            self.model = mnasnet0_5(pretrained=pretrained, num_classes=pretrained_num_classes)
        elif width_multiplier == 0.75:
            self.model = mnasnet0_75(pretrained=pretrained, num_classes=pretrained_num_classes)
        elif width_multiplier == 1.0:
            self.model = mnasnet1_0(pretrained=pretrained, num_classes=pretrained_num_classes)
        elif width_multiplier == 1.3:
            self.model = mnasnet1_3(pretrained=pretrained, num_classes=pretrained_num_classes)
        else:
            raise ValueError('no such value')
        self.init_weights(num_classes, pretrained_num_classes)

    def init_weights(self, num_classes, pretrained_num_classes):
        if num_classes != pretrained_num_classes:
            fc = self.model.classifier[1]
            fc_features = fc.in_features
            fc = nn.Linear(fc_features, num_classes)
            nn.init.normal_(fc.weight, 0, 0.01)
            nn.init.zeros_(fc.bias)
            self.model.classifier[1] = fc

    def train(self, mode: bool=True):
        super(TorchvisionMNASNet, self).train(mode=mode)
        if mode and (self.partial_bn or self.fix_bn):
            freezing_bn(self, partial_bn=self.partial_bn)
        return self

    def forward(self, x):
        x = self.model(x)
        return {KEY_OUTPUT: x}


class TorchvisionMobileNetV2(nn.Module, ABC):

    def __init__(self, num_classes=1000, width_multiplier=1.0, pretrained=False, pretrained_num_classes=1000, fix_bn=False, partial_bn=False, norm_layer=None):
        super(TorchvisionMobileNetV2, self).__init__()
        self.fix_bn = fix_bn
        self.partial_bn = partial_bn
        self.model = mobilenet_v2(pretrained=pretrained, width_mult=width_multiplier, norm_layer=norm_layer, num_classes=pretrained_num_classes)
        self.init_weights(num_classes, pretrained_num_classes)

    def init_weights(self, num_classes, pretrained_num_classes):
        if num_classes != pretrained_num_classes:
            fc = self.model.classifier[1]
            fc_features = fc.in_features
            fc = nn.Linear(fc_features, num_classes)
            nn.init.normal_(fc.weight, 0, 0.01)
            nn.init.zeros_(fc.bias)
            self.model.classifier[1] = fc

    def train(self, mode: bool=True):
        super(TorchvisionMobileNetV2, self).train(mode=mode)
        if mode and (self.partial_bn or self.fix_bn):
            freezing_bn(self, partial_bn=self.partial_bn)
        return self

    def forward(self, x):
        x = self.model(x)
        return {KEY_OUTPUT: x}


class TorchvisionMobileNetV3(nn.Module, ABC):

    def __init__(self, arch='mobilenet_v3_large', width_multiplier=1.0, num_classes=1000, pretrained=False, pretrained_num_classes=1000, fix_bn=False, partial_bn=False, norm_layer=None):
        super(TorchvisionMobileNetV3, self).__init__()
        self.fix_bn = fix_bn
        self.partial_bn = partial_bn
        if arch == 'mobilenet_v3_large':
            self.model = mobilenet_v3_large(pretrained=pretrained, _width_mult=width_multiplier, norm_layer=norm_layer, num_classes=pretrained_num_classes)
        elif arch == 'mobilenet_v3_small':
            self.model = mobilenet_v3_small(pretrained=pretrained, _width_mult=width_multiplier, norm_layer=norm_layer, num_classes=pretrained_num_classes)
        else:
            raise ValueError(f"{arch} doesn't exists")
        self.init_weights(num_classes, pretrained_num_classes)

    def init_weights(self, num_classes, pretrained_num_classes):
        if num_classes != pretrained_num_classes:
            fc = self.model.classifier[1]
            fc_features = fc.in_features
            fc = nn.Linear(fc_features, num_classes)
            nn.init.normal_(fc.weight, 0, 0.01)
            nn.init.zeros_(fc.bias)
            self.model.classifier[1] = fc

    def train(self, mode: bool=True):
        super(TorchvisionMobileNetV3, self).train(mode=mode)
        if mode and (self.partial_bn or self.fix_bn):
            freezing_bn(self, partial_bn=self.partial_bn)
        return self

    def forward(self, x):
        x = self.model(x)
        return {KEY_OUTPUT: x}


url_map = {'ghostnet_x1_0': 'https://github.com/ZJCV/ZCls/releases/download/v0.14.0/ghostnet_x1_0_imagenet_0000e4da.pth'}


class ResNet(BaseRecognizer):

    def __init__(self, cfg):
        super().__init__(cfg)

    def _init_weights(self, cfg):
        pretrained_local = cfg.MODEL.RECOGNIZER.PRETRAINED_LOCAL
        pretrained_num_classes = cfg.MODEL.RECOGNIZER.PRETRAINED_NUM_CLASSES
        num_classes = cfg.MODEL.HEAD.NUM_CLASSES
        model_name = cfg.MODEL.BACKBONE.ARCH
        assert isinstance(model_name, str)
        if model_name.startswith('repvgg'):
            if len(cfg.MODEL.CONV.ADD_BLOCKS) == 0:
                model_name += '_infer'
            else:
                model_name += '_train'
        load_pretrained_weights(self, model_name, weights_path=None if pretrained_local == '' else pretrained_local, load_fc=pretrained_num_classes == num_classes, verbose=True, url_map=url_map if cfg.MODEL.RECOGNIZER.PRETRAINED_REMOTE else None)


class OfficialResNeSt(nn.Module, ABC):

    def __init__(self, cfg, arch='resnest50_2s2x40d', dropout_rate=0.0, num_classes=1000, fix_bn=False, partial_bn=False, pretrained='', pretrained_num_classes=1000):
        super(OfficialResNeSt, self).__init__()
        self.num_classes = num_classes
        self.fix_bn = fix_bn
        self.partial_bn = partial_bn
        if arch == 'resnest50':
            self.model = resnest50(num_classes=pretrained_num_classes, final_drop=dropout_rate)
        elif arch == 'resnest50_2s2x40d':
            radix = 2
            groups = 2
            width_per_group = 40
            avg_first = False
            self.model = ResNet(Bottleneck, [3, 4, 6, 3], radix=radix, groups=groups, bottleneck_width=width_per_group, deep_stem=True, stem_width=32, avg_down=True, avd=True, avd_first=avg_first, final_drop=dropout_rate, num_classes=pretrained_num_classes)
        elif arch == 'resnest50_fast_2s2x40d':
            radix = 2
            groups = 2
            width_per_group = 40
            avg_first = True
            self.model = ResNet(Bottleneck, [3, 4, 6, 3], radix=radix, groups=groups, bottleneck_width=width_per_group, deep_stem=True, stem_width=32, avg_down=True, avd=True, avd_first=avg_first, final_drop=dropout_rate, num_classes=pretrained_num_classes)
        elif arch == 'resnest50_fast_2s1x64d':
            radix = 2
            groups = 1
            width_per_group = 64
            avg_first = True
            self.model = ResNet(Bottleneck, [3, 4, 6, 3], radix=radix, groups=groups, bottleneck_width=width_per_group, deep_stem=True, stem_width=32, avg_down=True, avd=True, avd_first=avg_first, final_drop=dropout_rate, num_classes=pretrained_num_classes)
        elif arch == 'resnest101':
            self.model = resnest101(num_classes=pretrained_num_classes, final_drop=dropout_rate)
        elif arch == 'resnest200':
            self.model = resnest200(num_classes=pretrained_num_classes, final_drop=dropout_rate)
        elif arch == 'resnest269':
            self.model = resnest269(num_classes=pretrained_num_classes, final_drop=dropout_rate)
        else:
            raise ValueError('no such value')
        self._init_weights(cfg)

    def _init_weights(self, cfg):
        pretrained_local = cfg.MODEL.RECOGNIZER.PRETRAINED_LOCAL
        pretrained_num_classes = cfg.MODEL.RECOGNIZER.PRETRAINED_NUM_CLASSES
        num_classes = cfg.MODEL.HEAD.NUM_CLASSES
        load_pretrained_weights(self, cfg.MODEL.BACKBONE.ARCH, weights_path=None if pretrained_local == '' else pretrained_local, load_fc=pretrained_num_classes == num_classes, verbose=True, url_map=None)

    def train(self, mode: bool=True) ->T:
        super(OfficialResNeSt, self).train(mode=mode)
        if mode and (self.partial_bn or self.fix_bn):
            freezing_bn(self, partial_bn=self.partial_bn)
        return self

    def forward(self, x):
        x = self.model(x)
        return {KEY_OUTPUT: x}


class TorchvisionResNet(nn.Module, ABC):

    def __init__(self, arch='resnet18', num_classes=1000, pretrained=False, pretrained_num_classes=1000, fix_bn=False, partial_bn=False, zero_init_residual=False):
        super(TorchvisionResNet, self).__init__()
        self.num_classes = num_classes
        self.fix_bn = fix_bn
        self.partial_bn = partial_bn
        if arch == 'resnet18':
            self.model = resnet18(pretrained=pretrained, num_classes=pretrained_num_classes, zero_init_residual=zero_init_residual)
        elif arch == 'resnet34':
            self.model = resnet34(pretrained=pretrained, num_classes=pretrained_num_classes, zero_init_residual=zero_init_residual)
        elif arch == 'resnet50':
            self.model = resnet50(pretrained=pretrained, num_classes=pretrained_num_classes, zero_init_residual=zero_init_residual)
        elif arch == 'resnet101':
            self.model = resnet101(pretrained=pretrained, num_classes=pretrained_num_classes, zero_init_residual=zero_init_residual)
        elif arch == 'resnet152':
            self.model = resnet152(pretrained=pretrained, num_classes=pretrained_num_classes, zero_init_residual=zero_init_residual)
        elif arch == 'resnext50_32x4d':
            self.model = resnext50_32x4d(pretrained=pretrained, num_classes=pretrained_num_classes, zero_init_residual=zero_init_residual)
        elif arch == 'resnext101_32x8d':
            self.model = resnext101_32x8d(pretrained=pretrained, num_classes=pretrained_num_classes, zero_init_residual=zero_init_residual)
        else:
            raise ValueError('no such value')
        self.init_weights(num_classes, pretrained_num_classes)

    def init_weights(self, num_classes, pretrained_num_classes):
        if num_classes != pretrained_num_classes:
            fc = self.model.fc
            fc_features = fc.in_features
            self.model.fc = nn.Linear(fc_features, num_classes)
            nn.init.normal_(self.model.fc.weight, 0, 0.01)
            nn.init.zeros_(self.model.fc.bias)

    def train(self, mode: bool=True) ->T:
        super(TorchvisionResNet, self).train(mode=mode)
        if mode and (self.partial_bn or self.fix_bn):
            freezing_bn(self, partial_bn=self.partial_bn)
        return self

    def forward(self, x):
        x = self.model(x)
        return {KEY_OUTPUT: x}


class TorchvisionShuffleNetV2(nn.Module, ABC):

    def __init__(self, arch='shufflenet_v2_x2_0', num_classes=1000, pretrained=False, pretrained_num_classes=1000, fix_bn=False, partial_bn=False):
        super(TorchvisionShuffleNetV2, self).__init__()
        self.fix_bn = fix_bn
        self.partial_bn = partial_bn
        if arch == 'shufflenet_v2_x2_0':
            self.model = shufflenet_v2_x2_0(pretrained=pretrained, num_classes=pretrained_num_classes)
        elif arch == 'shufflenet_v2_x1_5':
            self.model = shufflenet_v2_x1_5(pretrained=pretrained, num_classes=pretrained_num_classes)
        elif arch == 'shufflenet_v2_x1_0':
            self.model = shufflenet_v2_x1_0(pretrained=pretrained, num_classes=pretrained_num_classes)
        elif arch == 'shufflenet_v2_x0_5':
            self.model = shufflenet_v2_x0_5(pretrained=pretrained, num_classes=pretrained_num_classes)
        else:
            raise ValueError('no such value')
        self.init_weights(num_classes, pretrained_num_classes)

    def init_weights(self, num_classes, pretrained_num_classes):
        if num_classes != pretrained_num_classes:
            fc = self.model.fc
            fc_features = fc.in_features
            self.model.fc = nn.Linear(fc_features, num_classes)
            nn.init.normal_(self.model.fc.weight, 0, 0.01)
            nn.init.zeros_(self.model.fc.bias)

    def train(self, mode: bool=True) ->T:
        super(TorchvisionShuffleNetV2, self).train(mode=mode)
        if mode and (self.partial_bn or self.fix_bn):
            freezing_bn(self, partial_bn=self.partial_bn)
        return self

    def forward(self, x):
        x = self.model(x)
        return {KEY_OUTPUT: x}


import torch
from torch.nn import MSELoss, ReLU
from _paritybench_helpers import _mock_config, _mock_layer, _paritybench_base, _fails_compile


TESTCASES = [
    # (nn.Module, init_args, forward_args, jit_compiles)
    (BNAndPadLayer,
     lambda: ([], {'pad_pixels': 4, 'num_features': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (BasicBlock,
     lambda: ([], {'in_channels': 4, 'out_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (BatchNorm2d,
     lambda: ([], {'num_features': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (GhostBottleneck,
     lambda: ([], {'in_channels': 4, 'mid_channels': 4, 'out_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (GhostModule,
     lambda: ([], {'in_channels': 4, 'out_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (GroupNorm,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 32, 4, 4])], {}),
     True),
    (GroupNormWrapper,
     lambda: ([], {'num_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (IdentityBasedConv1x1,
     lambda: ([], {'channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (LayerNorm,
     lambda: ([], {'normalized_shape': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (MNASNetUint,
     lambda: ([], {'in_planes': 4, 'out_planes': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (MobileNetV1Block,
     lambda: ([], {'in_channels': 4, 'out_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (MobileNetV2Block,
     lambda: ([], {'in_channels': 4, 'out_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (MobileNetV2InvertedResidual,
     lambda: ([], {'in_channels': 4, 'out_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (MobileNetV3Head,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 960, 4, 4])], {}),
     True),
    (MobileNetV3Unit,
     lambda: ([], {'in_channels': 4, 'inner_channels': 4, 'out_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (NonLocal1DEmbeddedGaussian,
     lambda: ([], {'in_channels': 4}),
     lambda: ([torch.rand([4, 4, 4])], {}),
     False),
    (NonLocal2DEmbeddedGaussian,
     lambda: ([], {'in_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (NonLocal3DEmbeddedGaussian,
     lambda: ([], {'in_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4, 4])], {}),
     False),
    (PlaceHolder,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (RepVGGBackbone,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 64, 64])], {}),
     True),
    (ResNet3DBackbone,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 64, 64, 64])], {}),
     True),
    (ResNet3DBasicBlock,
     lambda: ([], {'in_planes': 4, 'out_planes': 4}),
     lambda: ([torch.rand([4, 4, 4, 4, 4])], {}),
     True),
    (SelectiveKernelConv2d,
     lambda: ([], {'in_channels': 4, 'out_channels': 4, 'stride': 1, 'groups': 1, 'reduction_rate': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (ShuffleNetV1Unit,
     lambda: ([], {'in_channels': 4, 'out_channels': 4, 'groups': 1, 'stride': 1}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (ShuffleNetV2Unit,
     lambda: ([], {'in_channels': 4, 'out_channels': 4, 'stride': 1}),
     lambda: ([torch.rand([4, 8, 4, 4])], {}),
     False),
    (SimplifiedNonLocal1DEmbeddedGaussian,
     lambda: ([], {'in_channels': 4}),
     lambda: ([torch.rand([4, 4, 4])], {}),
     False),
    (SimplifiedNonLocal2DEmbeddedGaussian,
     lambda: ([], {'in_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (SimplifiedNonLocal3DEmbeddedGaussian,
     lambda: ([], {'in_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4, 4])], {}),
     False),
    (SplitAttentionConv2d,
     lambda: ([], {'in_channels': 4, 'out_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (TorchvisionMNASNet,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 64, 64])], {}),
     False),
    (TorchvisionMobileNetV2,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 64, 64])], {}),
     False),
    (TorchvisionMobileNetV3,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 64, 64])], {}),
     False),
    (TorchvisionResNet,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 64, 64])], {}),
     False),
    (TorchvisionShuffleNetV2,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 64, 64])], {}),
     False),
    (rSoftMax,
     lambda: ([], {'radix': 4, 'cardinality': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
]

class Test_ZJCV_ZCls(_paritybench_base):
    def test_000(self):
        self._check(*TESTCASES[0])

    def test_001(self):
        self._check(*TESTCASES[1])

    def test_002(self):
        self._check(*TESTCASES[2])

    def test_003(self):
        self._check(*TESTCASES[3])

    def test_004(self):
        self._check(*TESTCASES[4])

    def test_005(self):
        self._check(*TESTCASES[5])

    def test_006(self):
        self._check(*TESTCASES[6])

    def test_007(self):
        self._check(*TESTCASES[7])

    def test_008(self):
        self._check(*TESTCASES[8])

    def test_009(self):
        self._check(*TESTCASES[9])

    def test_010(self):
        self._check(*TESTCASES[10])

    def test_011(self):
        self._check(*TESTCASES[11])

    def test_012(self):
        self._check(*TESTCASES[12])

    def test_013(self):
        self._check(*TESTCASES[13])

    def test_014(self):
        self._check(*TESTCASES[14])

    def test_015(self):
        self._check(*TESTCASES[15])

    def test_016(self):
        self._check(*TESTCASES[16])

    def test_017(self):
        self._check(*TESTCASES[17])

    def test_018(self):
        self._check(*TESTCASES[18])

    def test_019(self):
        self._check(*TESTCASES[19])

    def test_020(self):
        self._check(*TESTCASES[20])

    def test_021(self):
        self._check(*TESTCASES[21])

    def test_022(self):
        self._check(*TESTCASES[22])

    def test_023(self):
        self._check(*TESTCASES[23])

    def test_024(self):
        self._check(*TESTCASES[24])

    def test_025(self):
        self._check(*TESTCASES[25])

    def test_026(self):
        self._check(*TESTCASES[26])

    def test_027(self):
        self._check(*TESTCASES[27])

    def test_028(self):
        self._check(*TESTCASES[28])

    def test_029(self):
        self._check(*TESTCASES[29])

    def test_030(self):
        self._check(*TESTCASES[30])

    def test_031(self):
        self._check(*TESTCASES[31])

    def test_032(self):
        self._check(*TESTCASES[32])

    def test_033(self):
        self._check(*TESTCASES[33])

    def test_034(self):
        self._check(*TESTCASES[34])

