import sys
_module = sys.modules[__name__]
del sys
master = _module
chainer_ = _module
chainercv2 = _module
model_provider = _module
models = _module
airnet = _module
airnext = _module
alexnet = _module
alphapose_coco = _module
bagnet = _module
bamresnet = _module
bisenet = _module
bninception = _module
cbamresnet = _module
centernet = _module
cgnet = _module
channelnet = _module
common = _module
condensenet = _module
dabnet = _module
danet = _module
darknet = _module
darknet53 = _module
darts = _module
deeplabv3 = _module
densenet = _module
densenet_cifar = _module
diapreresnet = _module
diapreresnet_cifar = _module
diaresnet = _module
diaresnet_cifar = _module
dicenet = _module
diracnetv2 = _module
dla = _module
dpn = _module
drn = _module
efficientnet = _module
efficientnetedge = _module
espnetv2 = _module
fastscnn = _module
fastseresnet = _module
fbnet = _module
fcn8sd = _module
fdmobilenet = _module
fishnet = _module
fpenet = _module
ghostnet = _module
hardnet = _module
hrnet = _module
ibppose_coco = _module
icnet = _module
igcv3 = _module
inceptionresnetv1 = _module
inceptionresnetv2 = _module
inceptionv3 = _module
inceptionv4 = _module
irevnet = _module
jasper = _module
jasperdr = _module
lednet = _module
lffd = _module
lwopenpose_cmupan = _module
menet = _module
mixnet = _module
mnasnet = _module
mobilenet = _module
mobilenet_cub = _module
mobilenetb = _module
mobilenetv2 = _module
mobilenetv3 = _module
model_store = _module
nasnet = _module
nin_cifar = _module
ntsnet_cub = _module
nvpattexp = _module
octresnet = _module
others = _module
peleenet = _module
pnasnet = _module
polynet = _module
preresnet = _module
preresnet_cifar = _module
proxylessnas = _module
proxylessnas_cub = _module
pspnet = _module
pyramidnet = _module
pyramidnet_cifar = _module
quartznet = _module
regnet = _module
resattnet = _module
resdropresnet_cifar = _module
resnesta = _module
resnet = _module
resnet_cifar = _module
resnet_cub = _module
resneta = _module
resnetd = _module
resnext = _module
resnext_cifar = _module
rir_cifar = _module
ror_cifar = _module
scnet = _module
selecsls = _module
senet = _module
sepreresnet = _module
sepreresnet_cifar = _module
seresnet = _module
seresnet_cifar = _module
seresnet_cub = _module
seresnext = _module
shakedropresnet_cifar = _module
shakeshakeresnet_cifar = _module
sharesnet = _module
shufflenet = _module
shufflenetv2 = _module
shufflenetv2b = _module
simplepose_coco = _module
simpleposemobile_coco = _module
sinet = _module
sknet = _module
sparsenet = _module
spnasnet = _module
squeezenet = _module
squeezenext = _module
vgg = _module
voca = _module
vovnet = _module
wrn = _module
wrn1bit_cifar = _module
wrn_cifar = _module
xception = _module
xdensenet = _module
xdensenet_cifar = _module
zfnet = _module
dataset_utils = _module
datasets = _module
ade20k_seg_dataset = _module
cifar100_cls_dataset = _module
cifar10_cls_dataset = _module
cityscapes_seg_dataset = _module
coco_hpe1_dataset = _module
coco_hpe2_dataset = _module
coco_hpe3_dataset = _module
coco_seg_dataset = _module
cub200_2011_cls_dataset = _module
dataset_metainfo = _module
imagenet1k_cls_dataset = _module
seg_dataset = _module
svhn_cls_dataset = _module
voc_seg_dataset = _module
metrics = _module
cls_metrics = _module
det_metrics = _module
hpe_metrics = _module
metric = _module
seg_metrics = _module
seg_metrics_np = _module
setup = _module
utils = _module
env_stats = _module
logger_utils = _module
train_log_param_saver = _module
convert_models = _module
eval_ch = _module
eval_gl = _module
eval_gl_det = _module
eval_ke = _module
eval_pt = _module
eval_tf = _module
eval_tf2 = _module
convert_tf2_to_tfl = _module
demo_gl = _module
demo_pt = _module
demo_tf2 = _module
gluon = _module
asr_dataset = _module
coco_det_dataset = _module
coco_hpe2_dataset = _module
coco_hpe3_dataset = _module
hpatches_mch_dataset = _module
imagenet1k_rec_cls_dataset = _module
librispeech_asr_dataset = _module
mcv_asr_dataset = _module
widerface_det_dataset = _module
distillation = _module
gluoncv2 = _module
crunet = _module
crunetb = _module
diapreresnet = _module
diaresnet = _module
dla = _module
fractalnet_cifar = _module
ibnbresnet = _module
ibndensenet = _module
ibnresnet = _module
ibnresnext = _module
isqrtcovresnet = _module
msdnet = _module
octresnet_cifar = _module
quartznet = _module
regnetv = _module
res2net = _module
sinet = _module
superpointnet = _module
visemenet = _module
losses = _module
lr_scheduler = _module
asr_metrics = _module
seg_metrics_nd = _module
model_stats = _module
weighted_random_sampler = _module
keras_ = _module
kerascv = _module
load_model = _module
other = _module
cifar1 = _module
imagenet1k1 = _module
seg_utils1 = _module
top_k_accuracy1 = _module
train_ch_cifar = _module
train_ch_in1k = _module
eval_gl_mch = _module
eval_pt_mch = _module
khpa = _module
eval_gl_khpa = _module
khpa_cls_dataset = _module
khpa_utils = _module
train_gl_khpa = _module
pytorch = _module
cifar1 = _module
cub200_2011_utils1 = _module
imagenet1k1 = _module
seg_utils = _module
train_gl_seg = _module
prep_model = _module
dataset_utils = _module
asr_dataset = _module
cifar100_cls_dataset = _module
cifar10_cls_dataset = _module
coco_det_dataset = _module
coco_hpe1_dataset = _module
coco_hpe2_dataset = _module
coco_hpe3_dataset = _module
cub200_2011_cls_dataset = _module
hpatches_mch_dataset = _module
hpe_dataset = _module
imagenet1k_cls_dataset = _module
librispeech_asr_dataset = _module
mcv_asr_dataset = _module
mpii_hpe_dataset = _module
seg_dataset = _module
svhn_cls_dataset = _module
asr_metrics = _module
cls_metrics = _module
det_metrics = _module
hpe_metrics = _module
metric = _module
ret_metrics = _module
seg_metrics = _module
model_stats = _module
pytorchcv = _module
airnet = _module
airnext = _module
alexnet = _module
alphapose_coco = _module
bagnet = _module
bamresnet = _module
bisenet = _module
bninception = _module
cbamresnet = _module
centernet = _module
cgnet = _module
channelnet = _module
common = _module
condensenet = _module
contextnet = _module
dabnet = _module
danet = _module
darknet = _module
darknet53 = _module
darts = _module
deeplabv3 = _module
densenet = _module
densenet_cifar = _module
diapreresnet = _module
diapreresnet_cifar = _module
diaresnet = _module
diaresnet_cifar = _module
dicenet = _module
diracnetv2 = _module
dla = _module
dpn = _module
drn = _module
edanet = _module
efficientnet = _module
efficientnetedge = _module
enet = _module
erfnet = _module
esnet = _module
espcnet = _module
espnetv2 = _module
fastscnn = _module
fastseresnet = _module
fbnet = _module
fcn8sd = _module
fdmobilenet = _module
fishnet = _module
fpenet = _module
fractalnet_cifar = _module
ghostnet = _module
hardnet = _module
hrnet = _module
ibnbresnet = _module
ibndensenet = _module
ibnresnet = _module
ibnresnext = _module
ibppose_coco = _module
icnet = _module
igcv3 = _module
inceptionresnetv1 = _module
inceptionresnetv2 = _module
inceptionv3 = _module
inceptionv4 = _module
irevnet = _module
isqrtcovresnet = _module
jasper = _module
jasperdr = _module
lednet = _module
lffd = _module
linknet = _module
lwopenpose_cmupan = _module
menet = _module
mixnet = _module
mnasnet = _module
mobilenet = _module
mobilenet_cub = _module
mobilenetb = _module
mobilenetv2 = _module
mobilenetv3 = _module
model_store = _module
msdnet = _module
msdnet_cifar10 = _module
nasnet = _module
nin_cifar = _module
ntsnet_cub = _module
nvpattexp = _module
octresnet = _module
_espnet = _module
_inceptionresnetv1_ = _module
oth_espnet = _module
oth_inception_resnet_v1 = _module
oth_quartznet = _module
oth_vit = _module
peleenet = _module
pfpcnet = _module
pnasnet = _module
polynet = _module
preresnet = _module
preresnet_cifar = _module
prnet = _module
proxylessnas = _module
proxylessnas_cub = _module
pspnet = _module
pyramidnet = _module
pyramidnet_cifar = _module
quartznet = _module
regnet = _module
resattnet = _module
resdropresnet_cifar = _module
resnesta = _module
resnet = _module
resnet_cifar = _module
resnet_cub = _module
resneta = _module
resnetd = _module
resnext = _module
resnext_cifar = _module
revnet = _module
rir_cifar = _module
ror_cifar = _module
scnet = _module
segnet = _module
selecsls = _module
senet = _module
sepreresnet = _module
sepreresnet_cifar = _module
seresnet = _module
seresnet_cifar = _module
seresnet_cub = _module
seresnext = _module
shakedropresnet_cifar = _module
shakeshakeresnet_cifar = _module
sharesnet = _module
shufflenet = _module
shufflenetv2 = _module
shufflenetv2b = _module
simplepose_coco = _module
simpleposemobile_coco = _module
sinet = _module
sknet = _module
sparsenet = _module
spnasnet = _module
sqnet = _module
squeezenet = _module
squeezenext = _module
superpointnet = _module
tresnet = _module
unet = _module
vgg = _module
visemenet = _module
voca = _module
vovnet = _module
wrn = _module
wrn1bit_cifar = _module
wrn_cifar = _module
xception = _module
xdensenet = _module
xdensenet_cifar = _module
zfnet = _module
utils = _module
sotabench = _module
tensorflow2 = _module
cls_dataset = _module
coco_hpe2_dataset = _module
coco_hpe3_dataset = _module
tf2cv = _module
dla = _module
grmiposelite_coco = _module
quartznet = _module
sinet = _module
tensorflow_ = _module
tensorflowcv = _module
utils_tp = _module
tests = _module
convert_gl2pt_batchnorm = _module
convert_gl2pt_conv2d = _module
convert_gl2pt_dense = _module
convert_gl2tf2_avgpool2d = _module
convert_gl2tf2_batchnorm = _module
convert_gl2tf2_conv2d = _module
convert_gl2tf2_conv2d_b = _module
convert_gl2tf2_dwconv2d = _module
convert_gl2tf2_lstm = _module
convert_gl2tf_avgpool2d = _module
convert_gl2tf_batchnorm = _module
convert_gl2tf_conv1x1 = _module
convert_gl2tf_conv2d = _module
convert_gl2tf_dense = _module
convert_gl2tf_dwconv2d = _module
convert_gl2tf_gconv2d = _module
convert_gl2tf_maxpool2d = _module
train_ch = _module
train_gl = _module
train_gl_mealv2 = _module
train_ke = _module
train_pt = _module
train_tf = _module
train_tf2 = _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


from functools import partial


import math


import numpy as np


import logging


import re


import time


import torch


import pandas as pd


import torch.utils.data


import torchvision.transforms as transforms


import torchvision.datasets as datasets


import warnings


import random


import torch.nn as nn


import torch.backends.cudnn as cudnn


from torch.utils.data import DataLoader


from torch.utils.data.sampler import WeightedRandomSampler


import torch.utils.data as data


from torchvision.datasets import CIFAR100


from torchvision.datasets import CIFAR10


import copy


from torch.nn import functional as F


from torchvision.datasets import ImageFolder


from torchvision.datasets import SVHN


from collections import OrderedDict


from torch.autograd import Variable


import torch.nn.functional as F


import torch.nn.init as init


from inspect import isfunction


from torch.nn.parameter import Parameter


from torch.nn import init


from torch import nn


from tensorflow.keras.preprocessing.image import ImageDataGenerator


from tensorflow.keras.preprocessing.image import DirectoryIterator


import tensorflow as tf


import tensorflow.keras.layers as nn


def get_activation_layer(activation):
    """
    Create activation layer from string/function.

    Parameters:
    ----------
    activation : function, or str, or nn.Module
        Activation function or name of activation function.

    Returns:
    -------
    nn.Module
        Activation layer.
    """
    assert activation is not None
    if isfunction(activation):
        return activation()
    elif isinstance(activation, str):
        if activation == 'relu':
            return nn.ReLU(inplace=True)
        elif activation == 'relu6':
            return nn.ReLU6(inplace=True)
        elif activation == 'swish':
            return Swish()
        elif activation == 'hswish':
            return HSwish(inplace=True)
        elif activation == 'sigmoid':
            return nn.Sigmoid()
        elif activation == 'hsigmoid':
            return HSigmoid()
        elif activation == 'identity':
            return Identity()
        else:
            raise NotImplementedError()
    else:
        assert isinstance(activation, nn.Module)
        return activation


def conv1x1_block(in_channels, out_channels, stride=1, padding=0, groups=1, bias=False, use_bn=True, bn_eps=1e-05, activation=lambda : nn.ReLU(inplace=True)):
    """
    1x1 version of the standard convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int, or tuple/list of 2 int, or tuple/list of 4 int, default 0
        Padding value for convolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    """
    return ConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=padding, groups=groups, bias=bias, use_bn=use_bn, bn_eps=bn_eps, activation=activation)


def conv3x3_block(in_channels, out_channels, stride=1, padding=1, dilation=1, groups=1, bias=False, use_bn=True, bn_eps=1e-05, activation=lambda : nn.ReLU(inplace=True)):
    """
    3x3 version of the standard convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int, or tuple/list of 2 int, or tuple/list of 4 int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    """
    return ConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias, use_bn=use_bn, bn_eps=bn_eps, activation=activation)


def conv3x3(in_channels, out_channels, stride=1, padding=1, dilation=1, groups=1, bias=False):
    """
    Convolution 3x3 layer.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    """
    return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)


def conv1x1(in_channels, out_channels, stride=1, groups=1, bias=False):
    """
    Convolution 1x1 layer.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    """
    return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, groups=groups, bias=bias)


def conv7x7_block(in_channels, out_channels, stride=1, padding=3, dilation=1, groups=1, bias=False, use_bn=True, bn_eps=1e-05, activation=lambda : nn.ReLU(inplace=True)):
    """
    7x7 version of the standard convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    padding : int, or tuple/list of 2 int, or tuple/list of 4 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 3
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    """
    return ConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=7, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias, use_bn=use_bn, bn_eps=bn_eps, activation=activation)


def dwconv_block(in_channels, out_channels, kernel_size, stride=1, padding=1, dilation=1, bias=False, use_bn=True, bn_eps=1e-05, activation=lambda : nn.ReLU(inplace=True)):
    """
    Depthwise version of the standard convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    kernel_size : int or tuple/list of 2 int
        Convolution window size.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int, or tuple/list of 2 int, or tuple/list of 4 int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    """
    return ConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=out_channels, bias=bias, use_bn=use_bn, bn_eps=bn_eps, activation=activation)


def dwconv3x3_block(in_channels, out_channels, strides=1, padding=1, dilation=1, use_bias=False, bn_eps=1e-05, activation='relu', data_format='channels_last', **kwargs):
    """
    3x3 depthwise version of the standard convolution block (SINet version).

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    strides : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    use_bias : bool, default False
        Whether the layer uses a bias vector.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default 'relu'
        Activation function or name of activation function.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """
    return dwconv_block(in_channels=in_channels, out_channels=out_channels, kernel_size=3, strides=strides, padding=padding, dilation=dilation, use_bias=use_bias, bn_eps=bn_eps, activation=activation, data_format=data_format, **kwargs)


def is_channels_first():
    """
    Is tested data format channels first.

    Returns:
    -------
    bool
        A flag.
    """
    return K.image_data_format() == 'channels_first'


def round_channels(channels, divisor=8):
    """
    Round weighted channel number (make divisible operation).

    Parameters:
    ----------
    channels : int or float
        Original number of channels.
    divisor : int, default 8
        Alignment value.

    Returns:
    -------
    int
        Weighted number of channels.
    """
    rounded_channels = max(int(channels + divisor / 2.0) // divisor * divisor, divisor)
    if float(rounded_channels) < 0.9 * channels:
        rounded_channels += divisor
    return rounded_channels


class SEBlock(nn.Layer):
    """
    SINet version of Squeeze-and-Excitation block from 'Squeeze-and-Excitation Networks,'
    https://arxiv.org/abs/1709.01507.

    Parameters:
    ----------
    channels : int
        Number of channels.
    reduction : int, default 16
        Squeeze reduction value.
    round_mid : bool, default False
        Whether to round middle channel number (make divisible by 8).
    activation : function, or str, or nn.Module, default 'relu'
        Activation function after the first convolution.
    out_activation : function, or str, or nn.Module, default 'sigmoid'
        Activation function after the last convolution.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, channels, reduction=16, round_mid=False, mid_activation='relu', out_activation='sigmoid', data_format='channels_last', **kwargs):
        super(SEBlock, self).__init__(**kwargs)
        self.data_format = data_format
        self.use_conv2 = reduction > 1
        mid_channels = channels // reduction if not round_mid else round_channels(float(channels) / reduction)
        self.pool = nn.GlobalAveragePooling2D(data_format=data_format, name='pool')
        self.fc1 = nn.Dense(units=mid_channels, input_dim=channels, name='fc1')
        if self.use_conv2:
            self.activ = get_activation_layer(mid_activation, name='activ')
            self.fc2 = nn.Dense(units=channels, input_dim=mid_channels, name='fc2')
        self.sigmoid = get_activation_layer(out_activation, name='sigmoid')

    def call(self, x, training=None):
        w = self.pool(x)
        w = self.fc1(w)
        if self.use_conv2:
            w = self.activ(w)
            w = self.fc2(w)
        w = self.sigmoid(w)
        axis = -1 if is_channels_first(self.data_format) else 1
        w = tf.expand_dims(tf.expand_dims(w, axis=axis), axis=axis)
        x = x * w
        return x


def depthwise_conv3x3(channels, stride=1, padding=1, dilation=1, bias=False):
    """
    Depthwise convolution 3x3 layer.

    Parameters:
    ----------
    channels : int
        Number of input/output channels.
    strides : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    bias : bool, default False
        Whether the layer uses a bias vector.
    """
    return nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=3, stride=stride, padding=padding, dilation=dilation, groups=channels, bias=bias)


def channet_conv1x1(in_channels, out_channels, stride=1, groups=1, bias=False, dropout_rate=0.0, activate=True):
    """
    1x1 version of ChannelNet specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    dropout_rate : float, default 0.0
        Dropout rate.
    activate : bool, default True
        Whether activate the convolution block.
    """
    return ChannetConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=0, groups=groups, bias=bias, dropout_rate=dropout_rate, activate=activate)


def dwconv3x3(in_channels, out_channels, stride, bias=False):
    """
    3x3 depthwise version of the standard convolution layer.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int
        Strides of the convolution.
    bias : bool, default False
        Whether the layer uses a bias vector.
    """
    return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=1, groups=out_channels, bias=bias)


def channet_conv3x3(in_channels, out_channels, stride, padding=1, dilation=1, groups=1, bias=False, dropout_rate=0.0, activate=True):
    """
    3x3 version of the standard convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    dropout_rate : float, default 0.0
        Dropout rate.
    activate : bool, default True
        Whether activate the convolution block.
    """
    return ChannetConv(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias, dropout_rate=dropout_rate, activate=activate)


class DwsConvBlock(nn.Layer):
    """
    SINet version of depthwise separable convolution block with BatchNorms and activations at each convolution layers.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    kernel_size : int or tuple/list of 2 int
        Convolution window size.
    strides : int or tuple/list of 2 int
        Strides of the convolution.
    padding : int or tuple/list of 2 int
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    use_bias : bool, default False
        Whether the layer uses a bias vector.
    dw_use_bn : bool, default True
        Whether to use BatchNorm layer (depthwise convolution block).
    pw_use_bn : bool, default True
        Whether to use BatchNorm layer (pointwise convolution block).
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    dw_activation : function or str or None, default 'relu'
        Activation function after the depthwise convolution block.
    pw_activation : function or str or None, default 'relu'
        Activation function after the pointwise convolution block.
    se_reduction : int, default 0
        Squeeze reduction value (0 means no-se).
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, in_channels, out_channels, kernel_size, strides, padding, dilation=1, use_bias=False, dw_use_bn=True, pw_use_bn=True, bn_eps=1e-05, dw_activation='relu', pw_activation='relu', se_reduction=0, data_format='channels_last', **kwargs):
        super(DwsConvBlock, self).__init__(**kwargs)
        self.use_se = se_reduction > 0
        self.dw_conv = dwconv_block(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size, strides=strides, padding=padding, dilation=dilation, use_bias=use_bias, use_bn=dw_use_bn, bn_eps=bn_eps, activation=dw_activation, data_format=data_format, name='dw_conv')
        if self.use_se:
            self.se = SEBlock(channels=in_channels, reduction=se_reduction, round_mid=False, mid_activation=lambda : PReLU2(in_channels // se_reduction, data_format=data_format, name='activ'), out_activation=lambda : PReLU2(in_channels, data_format=data_format, name='sigmoid'), data_format=data_format, name='se')
        self.pw_conv = conv1x1_block(in_channels=in_channels, out_channels=out_channels, use_bias=use_bias, use_bn=pw_use_bn, bn_eps=bn_eps, activation=pw_activation, data_format=data_format, name='pw_conv')

    def call(self, x, training=None):
        x = self.dw_conv(x, training=None)
        if self.use_se:
            x = self.se(x, training=None)
        x = self.pw_conv(x, training=None)
        return x


def channel_shuffle(x, groups):
    """
    Channel shuffle operation from 'ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices,'
    https://arxiv.org/abs/1707.01083.

    Parameters:
    ----------
    x : Tensor
        Input tensor.
    groups : int
        Number of groups.

    Returns:
    -------
    Tensor
        Resulted tensor.
    """
    batch, channels, height, width = x.size()
    channels_per_group = channels // groups
    x = x.view(batch, groups, channels_per_group, height, width)
    x = torch.transpose(x, 1, 2).contiguous()
    x = x.view(batch, channels, height, width)
    return x


def channel_shuffle2(x, groups):
    """
    Channel shuffle operation from 'ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices,'
    https://arxiv.org/abs/1707.01083. The alternative version.

    Parameters:
    ----------
    x : Tensor
        Input tensor.
    groups : int
        Number of groups.

    Returns:
    -------
    Tensor
        Resulted tensor.
    """
    batch, channels, height, width = x.size()
    channels_per_group = channels // groups
    x = x.view(batch, channels_per_group, groups, height, width)
    x = torch.transpose(x, 1, 2).contiguous()
    x = x.view(batch, channels, height, width)
    return x


def condense_complex_conv1x1(in_channels, out_channels, groups):
    """
    1x1 version of the CondenseNet specific complex convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    groups : int
        Number of groups.
    """
    return CondenseComplexConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0, groups=groups)


def condense_simple_conv3x3(in_channels, out_channels, groups):
    """
    3x3 version of the CondenseNet specific simple convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    groups : int
        Number of groups.
    """
    return CondenseSimpleConv(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, groups=groups)


def dwsconv3x3_block(in_channels, out_channels, strides=1, padding=1, dilation=1, use_bias=False, dw_use_bn=True, pw_use_bn=True, bn_eps=1e-05, dw_activation='relu', pw_activation='relu', se_reduction=0, data_format='channels_last', **kwargs):
    """
    3x3 depthwise separable version of the standard convolution block (SINet version).

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    strides : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    use_bias : bool, default False
        Whether the layer uses a bias vector.
    dw_use_bn : bool, default True
        Whether to use BatchNorm layer (depthwise convolution block).
    pw_use_bn : bool, default True
        Whether to use BatchNorm layer (pointwise convolution block).
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    dw_activation : function or str or None, default 'relu'
        Activation function after the depthwise convolution block.
    pw_activation : function or str or None, default 'relu'
        Activation function after the pointwise convolution block.
    se_reduction : int, default 0
        Squeeze reduction value (0 means no-se).
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """
    return DwsConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=3, strides=strides, padding=padding, dilation=dilation, use_bias=use_bias, dw_use_bn=dw_use_bn, pw_use_bn=pw_use_bn, bn_eps=bn_eps, dw_activation=dw_activation, pw_activation=pw_activation, se_reduction=se_reduction, data_format=data_format, **kwargs)


def dwa_conv3x3_block(channels, stride=1, padding=1, dilation=1, bias=False, use_bn=True, bn_eps=1e-05, activation=lambda : nn.ReLU(inplace=True)):
    """
    3x3 version of the depthwise asymmetric separable convolution block.

    Parameters:
    ----------
    channels : int
        Number of input/output channels.
    stride : int, default 1
        Strides of the convolution.
    padding : int, default 1
        Padding value for convolution layer.
    dilation : int, default 1
        Dilation value for convolution layer.
    bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    """
    return DwaConvBlock(channels=channels, kernel_size=3, stride=stride, padding=padding, dilation=dilation, bias=bias, use_bn=use_bn, bn_eps=bn_eps, activation=activation)


def dark_convYxY(in_channels, out_channels, alpha, pointwise):
    """
    DarkNet unit.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    alpha : float
        Slope coefficient for Leaky ReLU activation.
    pointwise : bool
        Whether use 1x1 (pointwise) convolution or 3x3 convolution.
    """
    if pointwise:
        return conv1x1_block(in_channels=in_channels, out_channels=out_channels, activation=nn.LeakyReLU(negative_slope=alpha, inplace=True))
    else:
        return conv3x3_block(in_channels=in_channels, out_channels=out_channels, activation=nn.LeakyReLU(negative_slope=alpha, inplace=True))


def nasnet_batch_norm(channels):
    """
    NASNet specific Batch normalization layer.

    Parameters:
    ----------
    channels : int
        Number of channels in input data.
    """
    return nn.BatchNorm2d(num_features=channels, eps=0.001, momentum=0.1, affine=True)


def dws_branch_k3(in_channels, out_channels, stride=2, extra_padding=False, stem=False):
    """
    3x3 version of the PNASNet specific depthwise separable convolution branch.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 2
        Strides of the convolution.
    extra_padding : bool, default False
        Whether to use extra padding.
    stem : bool, default False
        Whether to use squeeze reduction if False.
    """
    return DwsBranch(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, extra_padding=extra_padding, stem=stem)


def dws_branch_k5(in_channels, out_channels, stride=2, extra_padding=False, stem=False):
    """
    5x5 version of the PNASNet specific depthwise separable convolution branch.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 2
        Strides of the convolution.
    extra_padding : bool, default False
        Whether to use extra padding.
    stem : bool, default False
        Whether to use squeeze reduction if False.
    """
    return DwsBranch(in_channels=in_channels, out_channels=out_channels, kernel_size=5, stride=stride, extra_padding=extra_padding, stem=stem)


def dws_branch_k7(in_channels, out_channels, stride=2, extra_padding=False):
    """
    7x7 version of the PNASNet specific depthwise separable convolution branch.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 2
        Strides of the convolution.
    extra_padding : bool, default False
        Whether to use extra padding.
    """
    return DwsBranch(in_channels=in_channels, out_channels=out_channels, kernel_size=7, stride=stride, extra_padding=extra_padding, stem=False)


def pnas_conv1x1(in_channels, out_channels, stride=1):
    """
    1x1 version of the PNASNet specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    """
    return NasConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=0, groups=1)


def darts_dws_branch3x3(channels, stride):
    """
    3x3 version of DARTS specific dilated convolution branch.

    Parameters:
    ----------
    channels : int
        Number of input/output channels.
    stride : int or tuple/list of 2 int
        Strides of the convolution.
    """
    return DartsDwsBranch(in_channels=channels, out_channels=channels, kernel_size=3, stride=stride, padding=1)


def darts_dws_conv3x3(channels, stride):
    """
    3x3 version of DARTS specific dilated convolution block.

    Parameters:
    ----------
    channels : int
        Number of input/output channels.
    stride : int or tuple/list of 2 int
        Strides of the convolution.
    """
    return DartsDwsConv(in_channels=channels, out_channels=channels, kernel_size=3, stride=stride, padding=2, dilation=2)


def darts_maxpool3x3(channels, stride):
    """
    DARTS specific 3x3 Max pooling layer.

    Parameters:
    ----------
    channels : int
        Number of input/output channels. Unused parameter.
    stride : int or tuple/list of 2 int
        Strides of the convolution.
    """
    assert channels > 0
    return nn.MaxPool2d(kernel_size=3, stride=stride, padding=1)


def darts_skip_connection(channels, stride):
    """
    DARTS specific skip connection layer.

    Parameters:
    ----------
    channels : int
        Number of input/output channels.
    stride : int or tuple/list of 2 int
        Strides of the convolution.
    """
    assert channels > 0
    if stride == 1:
        return Identity()
    else:
        assert stride == 2
        return DartsReduceBranch(in_channels=channels, out_channels=channels, stride=stride)


GENOTYPE_OPS = {'max_pool_3x3': darts_maxpool3x3, 'skip_connect': darts_skip_connection, 'dil_conv_3x3': darts_dws_conv3x3, 'sep_conv_3x3': darts_dws_branch3x3}


def darts_conv1x1(in_channels, out_channels, activate=True):
    """
    1x1 version of the DARTS specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    activate : bool, default True
        Whether activate the convolution block.
    """
    return DartsConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0, activate=activate)


class NasDualPathScheme(object):
    """
    NASNet specific scheme of dual path response for a module in a DualPathSequential module.

    Parameters:
    ----------
    can_skip_input : bool
        Whether can skip input for some modules.
    """

    def __init__(self, can_skip_input):
        super(NasDualPathScheme, self).__init__()
        self.can_skip_input = can_skip_input
    """
    Scheme function.

    Parameters:
    ----------
    module : nn.Module
        A module.
    x : Tensor
        Current processed tensor.
    x_prev : Tensor
        Previous processed tensor.

    Returns:
    -------
    x_next : Tensor
        Next processed tensor.
    x : Tensor
        Current processed tensor.
    """

    def __call__(self, module, x, x_prev):
        x_next = module(x, x_prev)
        if type(x_next) == tuple:
            x_next, x = x_next
        if self.can_skip_input and hasattr(module, 'skip_input') and module.skip_input:
            x = x_prev
        return x_next, x


def nasnet_dual_path_scheme_ordinal(module, x, _):
    """
    NASNet specific scheme of dual path response for an ordinal module with dual inputs/outputs in a DualPathSequential
    module.

    Parameters:
    ----------
    module : nn.Module
        A module.
    x : Tensor
        Current processed tensor.

    Returns:
    -------
    x_next : Tensor
        Next processed tensor.
    x : Tensor
        Current processed tensor.
    """
    return module(x), x


def nasnet_dual_path_sequential(return_two=True, first_ordinals=0, last_ordinals=0, can_skip_input=False):
    """
    NASNet specific dual path sequential container.

    Parameters:
    ----------
    return_two : bool, default True
        Whether to return two output after execution.
    first_ordinals : int, default 0
        Number of the first modules with single input/output.
    last_ordinals : int, default 0
        Number of the final modules with single input/output.
    dual_path_scheme : function
        Scheme of dual path response for a module.
    dual_path_scheme_ordinal : function
        Scheme of dual path response for an ordinal module.
    can_skip_input : bool, default False
        Whether can skip input for some modules.
    """
    return DualPathSequential(return_two=return_two, first_ordinals=first_ordinals, last_ordinals=last_ordinals, dual_path_scheme=NasDualPathScheme(can_skip_input=can_skip_input), dual_path_scheme_ordinal=nasnet_dual_path_scheme_ordinal)


def darts_conv3x3_s2(in_channels, out_channels, activate=True):
    """
    3x3 version of the DARTS specific convolution block with stride 2.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    activate : bool, default True
        Whether activate the convolution block.
    """
    return DartsConv(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=2, padding=1, activate=activate)


def stem2_unit(in_channels, out_channels):
    """
    DARTS Stem2 unit.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    """
    return darts_conv3x3_s2(in_channels=in_channels, out_channels=out_channels, activate=True)


def pre_conv1x1_block(in_channels, out_channels, stride=1, bias=False, use_bn=True, return_preact=False, activate=True):
    """
    1x1 version of the pre-activated convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    return_preact : bool, default False
        Whether return pre-activation.
    activate : bool, default True
        Whether activate the convolution block.
    """
    return PreConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=0, bias=bias, use_bn=use_bn, return_preact=return_preact, activate=activate)


def pre_conv3x3_block(in_channels, out_channels, stride=1, padding=1, dilation=1, bias=False, use_bn=True, return_preact=False, activate=True):
    """
    3x3 version of the pre-activated convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    return_preact : bool, default False
        Whether return pre-activation.
    activate : bool, default True
        Whether activate the convolution block.
    """
    return PreConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=padding, dilation=dilation, bias=bias, use_bn=use_bn, return_preact=return_preact, activate=activate)


def dirac_conv3x3(in_channels, out_channels):
    """
    3x3 version of the DiracNetV2 specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    """
    return DiracConv(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)


def get_channel_axis():
    """
    Get channel axis.

    Returns:
    -------
    int
        Channel axis.
    """
    return 1 if is_channels_first() else -1


class DLARoot(nn.Layer):
    """
    DLA root block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    residual : bool
        Whether use residual connection.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, in_channels, out_channels, residual, data_format='channels_last', **kwargs):
        super(DLARoot, self).__init__(**kwargs)
        self.residual = residual
        self.data_format = data_format
        self.conv = conv1x1_block(in_channels=in_channels, out_channels=out_channels, activation=None, data_format=data_format, name='conv')
        self.activ = nn.ReLU()

    def call(self, x2, x1, extra, training=None):
        last_branch = x2
        x = tf.concat([x2, x1] + list(extra), axis=get_channel_axis(self.data_format))
        x = self.conv(x, training=training)
        if self.residual:
            x += last_branch
        x = self.activ(x)
        return x


class DLATree(nn.Layer):
    """
    DLA tree unit. It's like iterative stage.

    Parameters:
    ----------
    levels : int
        Number of levels in the stage.
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    res_body_class : nn.Module
        Residual block body class.
    strides : int or tuple/list of 2 int
        Strides of the convolution in a residual block.
    root_residual : bool
        Whether use residual connection in the root.
    root_dim : int
        Number of input channels in the root block.
    first_tree : bool, default False
        Is this tree stage the first stage in the net.
    input_level : bool, default True
        Is this tree unit the first unit in the stage.
    return_down : bool, default False
        Whether return downsample result.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, levels, in_channels, out_channels, res_body_class, strides, root_residual, root_dim=0, first_tree=False, input_level=True, return_down=False, data_format='channels_last', **kwargs):
        super(DLATree, self).__init__(**kwargs)
        self.return_down = return_down
        self.add_down = input_level and not first_tree
        self.root_level = levels == 1
        if root_dim == 0:
            root_dim = 2 * out_channels
        if self.add_down:
            root_dim += in_channels
        if self.root_level:
            self.tree1 = DLAResBlock(in_channels=in_channels, out_channels=out_channels, strides=strides, body_class=res_body_class, return_down=True, data_format=data_format, name='tree1')
            self.tree2 = DLAResBlock(in_channels=out_channels, out_channels=out_channels, strides=1, body_class=res_body_class, return_down=False, data_format=data_format, name='tree2')
        else:
            self.tree1 = DLATree(levels=levels - 1, in_channels=in_channels, out_channels=out_channels, res_body_class=res_body_class, strides=strides, root_residual=root_residual, root_dim=0, input_level=False, return_down=True, data_format=data_format, name='tree1')
            self.tree2 = DLATree(levels=levels - 1, in_channels=out_channels, out_channels=out_channels, res_body_class=res_body_class, strides=1, root_residual=root_residual, root_dim=root_dim + out_channels, input_level=False, return_down=False, data_format=data_format, name='tree2')
        if self.root_level:
            self.root = DLARoot(in_channels=root_dim, out_channels=out_channels, residual=root_residual, data_format=data_format, name='root')

    def call(self, x, extra=None, training=None):
        extra = [] if extra is None else extra
        x1, down = self.tree1(x, training=training)
        if self.add_down:
            extra.append(down)
        if self.root_level:
            x2 = self.tree2(x1, training=training)
            x = self.root(x2, x1, extra, training=training)
        else:
            extra.append(x1)
            x = self.tree2(x1, extra, training=training)
        if self.return_down:
            return x, down
        else:
            return x


class DLAInitBlock(nn.Layer):
    """
    DLA specific initial block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, in_channels, out_channels, data_format='channels_last', **kwargs):
        super(DLAInitBlock, self).__init__(**kwargs)
        mid_channels = out_channels // 2
        self.conv1 = conv7x7_block(in_channels=in_channels, out_channels=mid_channels, data_format=data_format, name='conv1')
        self.conv2 = conv3x3_block(in_channels=mid_channels, out_channels=mid_channels, data_format=data_format, name='conv2')
        self.conv3 = conv3x3_block(in_channels=mid_channels, out_channels=out_channels, strides=2, data_format=data_format, name='conv3')

    def call(self, x, training=None):
        x = self.conv1(x, training=training)
        x = self.conv2(x, training=training)
        x = self.conv3(x, training=training)
        return x


def update_keras_shape(x):
    """
    Update Keras shape property.

    Parameters:
    ----------
    x : keras.backend tensor/variable/symbol
        Input tensor/variable/symbol.
    """
    if not hasattr(x, '_keras_shape'):
        x._keras_shape = tuple([(int(d) if d is not None and d != 0 else None) for d in x.shape])


def flatten(x, reshape=False):
    """
    Flattens the input to two dimensional.

    Parameters:
    ----------
    x : keras.backend tensor/variable/symbol
        Input tensor/variable/symbol.
    reshape : bool, default False
        Whether do reshape instead of flatten.

    Returns:
    -------
    keras.backend tensor/variable/symbol
        Resulted tensor/variable/symbol.
    """
    if not is_channels_first():

        def channels_last_flatten(z):
            z = K.permute_dimensions(z, pattern=(0, 3, 1, 2))
            z = K.reshape(z, shape=(-1, np.prod(K.int_shape(z)[1:])))
            update_keras_shape(z)
            return z
        return nn.Lambda(channels_last_flatten)(x)
    else:
        if reshape:
            x = nn.Reshape((-1,))(x)
        else:
            x = nn.Flatten()(x)
        return x


class DLA(tf.keras.Model):
    """
    DLA model from 'Deep Layer Aggregation,' https://arxiv.org/abs/1707.06484.

    Parameters:
    ----------
    levels : int
        Number of levels in each stage.
    channels : list of int
        Number of output channels for each stage.
    init_block_channels : int
        Number of output channels for the initial unit.
    res_body_class : nn.Module
        Residual block body class.
    residual_root : bool
        Whether use residual connection in the root blocks.
    in_channels : int, default 3
        Number of input channels.
    in_size : tuple of two ints, default (224, 224)
        Spatial size of the expected input image.
    classes : int, default 1000
        Number of classification classes.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, levels, channels, init_block_channels, res_body_class, residual_root, in_channels=3, in_size=(224, 224), classes=1000, data_format='channels_last', **kwargs):
        super(DLA, self).__init__(**kwargs)
        self.in_size = in_size
        self.classes = classes
        self.data_format = data_format
        self.features = SimpleSequential(name='features')
        self.features.add(DLAInitBlock(in_channels=in_channels, out_channels=init_block_channels, data_format=data_format, name='init_block'))
        in_channels = init_block_channels
        for i in range(len(levels)):
            levels_i = levels[i]
            out_channels = channels[i]
            first_tree = i == 0
            self.features.add(DLATree(levels=levels_i, in_channels=in_channels, out_channels=out_channels, res_body_class=res_body_class, strides=2, root_residual=residual_root, first_tree=first_tree, data_format=data_format, name='stage{}'.format(i + 1)))
            in_channels = out_channels
        self.features.add(nn.AveragePooling2D(pool_size=7, strides=1, data_format=data_format, name='final_pool'))
        self.output1 = conv1x1(in_channels=in_channels, out_channels=classes, use_bias=True, data_format=data_format, name='output1')

    def call(self, x, training=None):
        x = self.features(x, training=training)
        x = self.output1(x)
        x = flatten(x, self.data_format)
        return x


class BatchNorm(nn.BatchNormalization):
    """
    MXNet/Gluon-like batch normalization.

    Parameters:
    ----------
    momentum : float, default 0.9
        Momentum for the moving average.
    epsilon : float, default 1e-5
        Small float added to variance to avoid dividing by zero.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, momentum=0.9, epsilon=1e-05, data_format='channels_last', **kwargs):
        super(BatchNorm, self).__init__(axis=get_channel_axis(data_format), momentum=momentum, epsilon=epsilon, **kwargs)


class PreActivation(nn.Layer):
    """
    PreResNet like pure pre-activation block without convolution layer.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, in_channels, bn_eps=1e-05, data_format='channels_last', **kwargs):
        super(PreActivation, self).__init__(**kwargs)
        assert in_channels is not None
        self.bn = BatchNorm(epsilon=bn_eps, data_format=data_format, name='bn')
        self.activ = PReLU2(in_channels, data_format=data_format, name='activ')

    def call(self, x, training=None):
        x = self.bn(x, training=None)
        x = self.activ(x)
        return x


def dpn_batch_norm(channels):
    """
    DPN specific Batch normalization layer.

    Parameters:
    ----------
    channels : int
        Number of channels in input data.
    """
    return nn.BatchNorm2d(num_features=channels, eps=0.001)


def dpn_conv1x1(in_channels, out_channels, stride=1):
    """
    1x1 version of the DPN specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    """
    return DPNConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=0, groups=1)


def dpn_conv3x3(in_channels, out_channels, stride, groups):
    """
    3x3 version of the DPN specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int
        Strides of the convolution.
    groups : int
        Number of groups.
    """
    return DPNConv(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=1, groups=groups)


def drn_conv3x3(in_channels, out_channels, stride, dilation, activate):
    """
    3x3 version of the DRN specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int
        Strides of the convolution.
    dilation : int or tuple/list of 2 int
        Padding/dilation value for convolution layer.
    activate : bool
        Whether activate the convolution block.
    """
    return DRNConv(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=dilation, dilation=dilation, activate=activate)


def drn_conv1x1(in_channels, out_channels, stride, activate):
    """
    1x1 version of the DRN specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int
        Strides of the convolution.
    activate : bool
        Whether activate the convolution block.
    """
    return DRNConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=0, dilation=1, activate=activate)


def drn_init_block(in_channels, out_channels):
    """
    DRN specific initial block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    """
    return DRNConv(in_channels=in_channels, out_channels=out_channels, kernel_size=7, stride=1, padding=3, dilation=1, activate=True)


def asym_conv3x3_block(padding=1, **kwargs):
    """
    3x3 asymmetric separable convolution block.

    Parameters:
    ----------
    channels : int
        Number of input/output channels.
    padding : int, default 1
        Padding value for convolution layer.
    dilation : int, default 1
        Dilation value for convolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    lw_use_bn : bool, default True
        Whether to use BatchNorm layer (leftwise convolution block).
    rw_use_bn : bool, default True
        Whether to use BatchNorm layer (rightwise convolution block).
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    lw_activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function after the leftwise convolution block.
    rw_activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function after the rightwise convolution block.
    """
    return AsymConvBlock(kernel_size=3, padding=padding, **kwargs)


def calc_tf_padding(x, kernel_size, stride=1, dilation=1):
    """
    Calculate TF-same like padding size.

    Parameters:
    ----------
    x : tensor
        Input tensor.
    kernel_size : int
        Convolution window size.
    stride : int, default 1
        Strides of the convolution.
    dilation : int, default 1
        Dilation value for convolution layer.

    Returns:
    -------
    tuple of 4 int
        The size of the padding.
    """
    height, width = x.size()[2:]
    oh = math.ceil(float(height) / stride)
    ow = math.ceil(float(width) / stride)
    pad_h = max((oh - 1) * stride + (kernel_size - 1) * dilation + 1 - height, 0)
    pad_w = max((ow - 1) * stride + (kernel_size - 1) * dilation + 1 - width, 0)
    return pad_h // 2, pad_h - pad_h // 2, pad_w // 2, pad_w - pad_w // 2


def dwconv5x5_block(in_channels, out_channels, stride=1, padding=2, dilation=1, bias=False, bn_eps=1e-05, activation=lambda : nn.ReLU(inplace=True)):
    """
    5x5 depthwise version of the standard convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int, or tuple/list of 2 int, or tuple/list of 4 int, default 2
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    bias : bool, default False
        Whether the layer uses a bias vector.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    """
    return dwconv_block(in_channels=in_channels, out_channels=out_channels, kernel_size=5, stride=stride, padding=padding, dilation=dilation, bias=bias, bn_eps=bn_eps, activation=activation)


def deconv3x3_block(padding=1, out_padding=1, **kwargs):
    """
    3x3 version of the deconvolution block with batch normalization and activation.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int
        Strides of the deconvolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for deconvolution layer.
    ext_padding : tuple/list of 4 int, default None
        Extra padding value for deconvolution layer.
    out_padding : int or tuple/list of 2 int, default 1
        Output padding value for deconvolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for deconvolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    """
    return DeconvBlock(kernel_size=3, padding=padding, out_padding=out_padding, **kwargs)


class AvgPool2d(nn.Layer):
    """
    Average pooling operation for two dimensional (spatial) data.

    Parameters:
    ----------
    pool_size : int or tuple/list of 2 int
        Size of the max pooling windows.
    strides : int or tuple/list of 2 int
        Strides of the convolution.
    padding : int or tuple/list of 2 int
        Padding value for convolution layer.
    ceil_mode : bool, default False
        When `True`, will use ceil instead of floor to compute the output shape.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, pool_size, strides, padding=0, ceil_mode=False, data_format='channels_last', **kwargs):
        super(AvgPool2d, self).__init__(**kwargs)
        if isinstance(pool_size, int):
            pool_size = pool_size, pool_size
        if isinstance(strides, int):
            strides = strides, strides
        if isinstance(padding, int):
            padding = padding, padding
        self.use_stride = strides[0] > 1 or strides[1] > 1
        self.ceil_mode = ceil_mode and self.use_stride
        self.use_pad = padding[0] > 0 or padding[1] > 0
        if self.ceil_mode:
            self.padding = padding
            self.pool_size = pool_size
            self.strides = strides
            self.data_format = data_format
        elif self.use_pad:
            if is_channels_first(data_format):
                self.paddings_tf = [[0, 0], [0, 0], [padding[0]] * 2, [padding[1]] * 2]
            else:
                self.paddings_tf = [[0, 0], [padding[0]] * 2, [padding[1]] * 2, [0, 0]]
        self.pool = nn.AveragePooling2D(pool_size=pool_size, strides=1, padding='valid', data_format=data_format, name='pool')
        if self.use_stride:
            self.stride_pool = nn.AveragePooling2D(pool_size=1, strides=strides, padding='valid', data_format=data_format, name='stride_pool')

    def call(self, x, training=None):
        if self.ceil_mode:
            x_shape = x.get_shape().as_list()
            if is_channels_first(self.data_format):
                height = x_shape[2]
                width = x_shape[3]
            else:
                height = x_shape[1]
                width = x_shape[2]
            padding = self.padding
            out_height = float(height + 2 * padding[0] - self.pool_size[0]) / self.strides[0] + 1.0
            out_width = float(width + 2 * padding[1] - self.pool_size[1]) / self.strides[1] + 1.0
            if math.ceil(out_height) > math.floor(out_height):
                padding = padding[0] + 1, padding[1]
            if math.ceil(out_width) > math.floor(out_width):
                padding = padding[0], padding[1] + 1
            if padding[0] > 0 or padding[1] > 0:
                if is_channels_first(self.data_format):
                    paddings_tf = [[0, 0], [0, 0], [padding[0]] * 2, [padding[1]] * 2]
                else:
                    paddings_tf = [[0, 0], [padding[0]] * 2, [padding[1]] * 2, [0, 0]]
                x = tf.pad(x, paddings=paddings_tf)
        elif self.use_pad:
            x = tf.pad(x, paddings=self.paddings_tf)
        x = self.pool(x)
        if self.use_stride:
            x = self.stride_pool(x)
        return x


class FDWConvBlock(nn.Layer):
    """
    Factorized depthwise separable convolution block with BatchNorms and activations at each convolution layers.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    kernel_size : int
        Convolution window size.
    strides : int or tuple/list of 2 int
        Strides of the convolution.
    padding : int
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    use_bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default 'relu'
        Activation function after the each convolution block.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, in_channels, out_channels, kernel_size, strides, padding, dilation=1, use_bias=False, use_bn=True, bn_eps=1e-05, activation='relu', data_format='channels_last', **kwargs):
        super(FDWConvBlock, self).__init__(**kwargs)
        assert use_bn
        self.activate = activation is not None
        self.v_conv = dwconv_block(in_channels=in_channels, out_channels=out_channels, kernel_size=(kernel_size, 1), strides=strides, padding=(padding, 0), dilation=dilation, use_bias=use_bias, use_bn=use_bn, bn_eps=bn_eps, activation=None, data_format=data_format, name='v_conv')
        self.h_conv = dwconv_block(in_channels=in_channels, out_channels=out_channels, kernel_size=(1, kernel_size), strides=strides, padding=(0, padding), dilation=dilation, use_bias=use_bias, use_bn=use_bn, bn_eps=bn_eps, activation=None, data_format=data_format, name='h_conv')
        if self.activate:
            self.act = get_activation_layer(activation, name='act')

    def call(self, x, training=None):
        x = self.v_conv(x, training=None) + self.h_conv(x, training=None)
        if self.activate:
            x = self.act(x)
        return x


def fdwconv3x3_block(in_channels, out_channels, strides=1, padding=1, dilation=1, use_bias=False, use_bn=True, bn_eps=1e-05, activation='relu', data_format='channels_last', **kwargs):
    """
    3x3 factorized depthwise version of the standard convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    strides : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    use_bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default 'relu'
        Activation function or name of activation function.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """
    return FDWConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=3, strides=strides, padding=padding, dilation=dilation, use_bias=use_bias, use_bn=use_bn, bn_eps=bn_eps, activation=activation, data_format=data_format, **kwargs)


def fdwconv5x5_block(in_channels, out_channels, strides=1, padding=2, dilation=1, use_bias=False, use_bn=True, bn_eps=1e-05, activation='relu', data_format='channels_last', **kwargs):
    """
    5x5 factorized depthwise version of the standard convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    strides : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    use_bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default 'relu'
        Activation function or name of activation function.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """
    return FDWConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=5, strides=strides, padding=padding, dilation=dilation, use_bias=use_bias, use_bn=use_bn, bn_eps=bn_eps, activation=activation, data_format=data_format, **kwargs)


class SBBlock(nn.Layer):
    """
    SB-block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    kernel_size : int
        Convolution window size for a factorized depthwise separable convolution block.
    scale_factor : int
        Scale factor.
    size : tuple of 2 int
        Spatial size of the output tensor for the bilinear upsampling operation.
    bn_eps : float
        Small float added to variance in Batch norm.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, in_channels, out_channels, kernel_size, scale_factor, size, bn_eps, data_format='channels_last', **kwargs):
        super(SBBlock, self).__init__(**kwargs)
        self.use_scale = scale_factor > 1
        if self.use_scale:
            self.down_scale = AvgPool2d(pool_size=scale_factor, strides=scale_factor, data_format=data_format, name='down_scale')
            self.up_scale = InterpolationBlock(scale_factor=scale_factor, out_size=size, data_format=data_format, name='up_scale')
        use_fdw = scale_factor > 0
        if use_fdw:
            fdwconv3x3_class = fdwconv3x3_block if kernel_size == 3 else fdwconv5x5_block
            self.conv1 = fdwconv3x3_class(in_channels=in_channels, out_channels=in_channels, bn_eps=bn_eps, activation=lambda : PReLU2(in_channels, data_format=data_format, name='activ'), data_format=data_format, name='conv1')
        else:
            self.conv1 = dwconv3x3_block(in_channels=in_channels, out_channels=in_channels, bn_eps=bn_eps, activation=lambda : PReLU2(in_channels, data_format=data_format, name='activ'), data_format=data_format, name='conv1')
        self.conv2 = conv1x1(in_channels=in_channels, out_channels=out_channels, data_format=data_format, name='conv2')
        self.bn = BatchNorm(epsilon=bn_eps, data_format=data_format, name='bn')

    def call(self, x, training=None):
        if self.use_scale:
            x = self.down_scale(x)
        x = self.conv1(x, training=None)
        x = self.conv2(x, training=None)
        if self.use_scale:
            x = self.up_scale(x)
        x = self.bn(x, training=None)
        return x


class ESPBlock(nn.Layer):
    """
    ESP block, which is based on the following principle: Reduce ---> Split ---> Transform --> Merge.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    kernel_sizes : list of int
        Convolution window size for branches.
    scale_factors : list of int
        Scale factor for branches.
    use_residual : bool
        Whether to use residual connection.
    in_size : tuple of 2 int
        Spatial size of the output tensor for the bilinear upsampling operation.
    bn_eps : float
        Small float added to variance in Batch norm.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, in_channels, out_channels, kernel_sizes, scale_factors, use_residual, in_size, bn_eps, data_format='channels_last', **kwargs):
        super(ESPBlock, self).__init__(**kwargs)
        self.use_residual = use_residual
        groups = len(kernel_sizes)
        mid_channels = int(out_channels / groups)
        res_channels = out_channels - groups * mid_channels
        self.conv = conv1x1(in_channels=in_channels, out_channels=mid_channels, groups=groups, data_format=data_format, name='conv')
        self.c_shuffle = ChannelShuffle(channels=mid_channels, groups=groups, data_format=data_format, name='c_shuffle')
        self.branches = Concurrent(data_format=data_format, name='branches')
        for i in range(groups):
            out_channels_i = mid_channels + res_channels if i == 0 else mid_channels
            self.branches.add(SBBlock(in_channels=mid_channels, out_channels=out_channels_i, kernel_size=kernel_sizes[i], scale_factor=scale_factors[i], size=in_size, bn_eps=bn_eps, data_format=data_format, name='branch{}'.format(i + 1)))
        self.preactiv = PreActivation(in_channels=out_channels, bn_eps=bn_eps, data_format=data_format, name='preactiv')

    def call(self, x, training=None):
        if self.use_residual:
            identity = x
        x = self.conv(x)
        x = self.c_shuffle(x)
        x = self.branches(x, training=None)
        if self.use_residual:
            x = identity + x
        x = self.preactiv(x, training=None)
        return x


def channel_squeeze(x, groups):
    """
    Channel squeeze operation.

    Parameters:
    ----------
    x : Tensor
        Input tensor.
    groups : int
        Number of groups.

    Returns:
    -------
    Tensor
        Resulted tensor.
    """
    batch, channels, height, width = x.size()
    channels_per_group = channels // groups
    x = x.view(batch, channels_per_group, groups, height, width).sum(dim=2)
    return x


def drop_conv3x3_block(in_channels, out_channels, stride=1, padding=1, bias=False, dropout_prob=0.0):
    """
    3x3 version of the convolution block with dropout.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    bias : bool, default False
        Whether the layer uses a bias vector.
    dropout_rate : float, default 0.0
        Parameter of Dropout layer. Faction of the input units to drop.
    """
    return DropConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=padding, bias=bias, dropout_prob=dropout_prob)


def invdwsconv3x3_block(in_channels, out_channels, stride=1, padding=1, dilation=1, bias=False, bn_eps=1e-05, pw_activation=lambda : nn.ReLU(inplace=True), dw_activation=lambda : nn.ReLU(inplace=True)):
    """
    3x3 inverse depthwise separable version of the standard convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    bias : bool, default False
        Whether the layer uses a bias vector.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    pw_activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function after the pointwise convolution block.
    dw_activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function after the depthwise convolution block.
    """
    return InvDwsConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=padding, dilation=dilation, bias=bias, bn_eps=bn_eps, pw_activation=pw_activation, dw_activation=dw_activation)


def ibnb_conv7x7_block(in_channels, out_channels, stride=1, padding=3, bias=False, activate=True):
    """
    7x7 version of the IBN(b)-ResNet specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 3
        Padding value for convolution layer.
    bias : bool, default False
        Whether the layer uses a bias vector.
    activate : bool, default True
        Whether activate the convolution block.
    """
    return IBNbConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=7, stride=stride, padding=padding, bias=bias, activate=activate)


def ibn_pre_conv1x1_block(in_channels, out_channels, stride=1, use_ibn=False, return_preact=False):
    """
    1x1 version of the IBN-Net specific pre-activated convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    use_ibn : bool, default False
        Whether use Instance-Batch Normalization.
    return_preact : bool, default False
        Whether return pre-activation.
    """
    return IBNPreConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=0, use_ibn=use_ibn, return_preact=return_preact)


def ibn_conv1x1_block(in_channels, out_channels, stride=1, groups=1, bias=False, use_ibn=False, activate=True):
    """
    1x1 version of the IBN-Net specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    use_ibn : bool, default False
        Whether use Instance-Batch Normalization.
    activate : bool, default True
        Whether activate the convolution block.
    """
    return IBNConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=0, groups=groups, bias=bias, use_ibn=use_ibn, activate=activate)


class CovPool(torch.autograd.Function):
    """
    Covariance pooling function.
    """

    @staticmethod
    def forward(ctx, x):
        batch, channels, height, width = x.size()
        n = height * width
        xn = x.reshape(batch, channels, n)
        identity_bar = (1.0 / n * torch.eye(n, dtype=xn.dtype, device=xn.device)).unsqueeze(dim=0).repeat(batch, 1, 1)
        ones_bar = torch.full((batch, n, n), fill_value=-1.0 / n / n, dtype=xn.dtype, device=xn.device)
        i_bar = identity_bar + ones_bar
        sigma = xn.bmm(i_bar).bmm(xn.transpose(1, 2))
        ctx.save_for_backward(x, i_bar)
        return sigma

    @staticmethod
    def backward(ctx, grad_sigma):
        x, i_bar = ctx.saved_tensors
        batch, channels, height, width = x.size()
        n = height * width
        xn = x.reshape(batch, channels, n)
        grad_x = grad_sigma + grad_sigma.transpose(1, 2)
        grad_x = grad_x.bmm(xn).bmm(i_bar)
        grad_x = grad_x.reshape(batch, channels, height, width)
        return grad_x


class NewtonSchulzSqrt(torch.autograd.Function):
    """
    Newton-Schulz iterative matrix square root function.

    Parameters:
    ----------
    x : Tensor
        Input tensor (batch * cols * rows).
    n : int
        Number of iterations (n > 1).
    """

    @staticmethod
    def forward(ctx, x, n):
        assert n > 1
        batch, cols, rows = x.size()
        assert cols == rows
        m = cols
        identity = torch.eye(m, dtype=x.dtype, device=x.device).unsqueeze(dim=0).repeat(batch, 1, 1)
        x_trace = (x * identity).sum(dim=(1, 2), keepdim=True)
        a = x / x_trace
        i3 = 3.0 * identity
        yi = torch.zeros(batch, n - 1, m, m, dtype=x.dtype, device=x.device)
        zi = torch.zeros(batch, n - 1, m, m, dtype=x.dtype, device=x.device)
        b2 = 0.5 * (i3 - a)
        yi[:, 0, :, :] = a.bmm(b2)
        zi[:, 0, :, :] = b2
        for i in range(1, n - 1):
            b2 = 0.5 * (i3 - zi[:, i - 1, :, :].bmm(yi[:, i - 1, :, :]))
            yi[:, i, :, :] = yi[:, i - 1, :, :].bmm(b2)
            zi[:, i, :, :] = b2.bmm(zi[:, i - 1, :, :])
        b2 = 0.5 * (i3 - zi[:, n - 2, :, :].bmm(yi[:, n - 2, :, :]))
        yn = yi[:, n - 2, :, :].bmm(b2)
        x_trace_sqrt = torch.sqrt(x_trace)
        c = yn * x_trace_sqrt
        ctx.save_for_backward(x, x_trace, a, yi, zi, yn, x_trace_sqrt)
        ctx.n = n
        return c

    @staticmethod
    def backward(ctx, grad_c):
        x, x_trace, a, yi, zi, yn, x_trace_sqrt = ctx.saved_tensors
        n = ctx.n
        batch, m, _ = x.size()
        identity0 = torch.eye(m, dtype=x.dtype, device=x.device)
        identity = identity0.unsqueeze(dim=0).repeat(batch, 1, 1)
        i3 = 3.0 * identity
        grad_yn = grad_c * x_trace_sqrt
        b = i3 - yi[:, n - 2, :, :].bmm(zi[:, n - 2, :, :])
        grad_yi = 0.5 * (grad_yn.bmm(b) - zi[:, n - 2, :, :].bmm(yi[:, n - 2, :, :]).bmm(grad_yn))
        grad_zi = -0.5 * yi[:, n - 2, :, :].bmm(grad_yn).bmm(yi[:, n - 2, :, :])
        for i in range(n - 3, -1, -1):
            b = i3 - yi[:, i, :, :].bmm(zi[:, i, :, :])
            ziyi = zi[:, i, :, :].bmm(yi[:, i, :, :])
            grad_yi_m1 = 0.5 * (grad_yi.bmm(b) - zi[:, i, :, :].bmm(grad_zi).bmm(zi[:, i, :, :]) - ziyi.bmm(grad_yi))
            grad_zi_m1 = 0.5 * (b.bmm(grad_zi) - yi[:, i, :, :].bmm(grad_yi).bmm(yi[:, i, :, :]) - grad_zi.bmm(ziyi))
            grad_yi = grad_yi_m1
            grad_zi = grad_zi_m1
        grad_a = 0.5 * (grad_yi.bmm(i3 - a) - grad_zi - a.bmm(grad_yi))
        x_trace_sqr = x_trace * x_trace
        grad_atx_trace = (grad_a.transpose(1, 2).bmm(x) * identity).sum(dim=(1, 2), keepdim=True)
        grad_cty_trace = (grad_c.transpose(1, 2).bmm(yn) * identity).sum(dim=(1, 2), keepdim=True)
        grad_x_extra = (0.5 * grad_cty_trace / x_trace_sqrt - grad_atx_trace / x_trace_sqr).repeat(1, m, m) * identity
        grad_x = grad_a / x_trace + grad_x_extra
        return grad_x, None


class Triuvec(torch.autograd.Function):
    """
    Extract upper triangular part of matrix into vector form.
    """

    @staticmethod
    def forward(ctx, x):
        batch, cols, rows = x.size()
        assert cols == rows
        n = cols
        triuvec_inds = torch.ones(n, n).triu().view(n * n).nonzero()
        x_vec = x.reshape(batch, -1)
        y = x_vec[:, triuvec_inds]
        ctx.save_for_backward(x, triuvec_inds)
        return y

    @staticmethod
    def backward(ctx, grad_y):
        x, triuvec_inds = ctx.saved_tensors
        batch, n, _ = x.size()
        grad_x = torch.zeros_like(x).view(batch, -1)
        grad_x[:, triuvec_inds] = grad_y
        grad_x = grad_x.view(batch, n, n)
        return grad_x


def outmask_fill(x, x_len, value=0.0):
    """
    Masked fill a tensor.

    Parameters:
    ----------
    x : tensor
        Input tensor.
    x_len : tensor
        Tensor with lengths.
    value : float, default 0.0
        Filled value.

    Returns:
    -------
    tensor
        Resulted tensor.
    """
    max_len = x.size(2)
    mask = torch.arange(max_len).expand(len(x_len), max_len) >= x_len.unsqueeze(1)
    mask = mask.unsqueeze(dim=1)
    x = x.masked_fill(mask=mask, value=value)
    return x


def masked_normalize2(x, x_len):
    """
    Normalize a tensor with mask (scheme #2).

    Parameters:
    ----------
    x : tensor
        Input tensor.
    x_len : tensor
        Tensor with lengths.

    Returns:
    -------
    tensor
        Resulted tensor.
    """
    x = outmask_fill(x, x_len)
    x_mean = x.sum(dim=2) / x_len.unsqueeze(dim=1)
    x2_mean = x.square().sum(dim=2) / x_len.unsqueeze(dim=1)
    x_std = (x2_mean - x_mean.square()).sqrt()
    x = (x - x_mean.unsqueeze(dim=2)) / x_std.unsqueeze(dim=2)
    return x


def mask_conv1d1(in_channels, out_channels, stride=1, groups=1, bias=False):
    """
    Masked 1-dim kernel version of the 1D convolution layer.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int, default 1
        Strides of the convolution.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    """
    return MaskConv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, groups=groups, bias=bias)


def mask_conv1d1_block(in_channels, out_channels, stride=1, padding=0, **kwargs):
    """
    1-dim kernel version of the masked 1D convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int, default 1
        Strides of the convolution.
    padding : int, default 0
        Padding value for convolution layer.
    """
    return MaskConvBlock1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=padding, **kwargs)


class CtcDecoder(object):
    """
    CTC decoder (to decode a sequence of labels to words).

    Parameters:
    ----------
    vocabulary : list of str
        Vocabulary of the dataset.
    """

    def __init__(self, vocabulary):
        super().__init__()
        self.blank_id = len(vocabulary)
        self.labels_map = dict([(i, vocabulary[i]) for i in range(len(vocabulary))])

    def __call__(self, predictions):
        """
        Decode a sequence of labels to words.

        Parameters:
        ----------
        predictions : np.array of int or list of list of int
            Tensor with predicted labels.

        Returns:
        -------
        list of str
            Words.
        """
        hypotheses = []
        for prediction in predictions:
            decoded_prediction = []
            previous = self.blank_id
            for p in prediction:
                if (p != previous or previous == self.blank_id) and p != self.blank_id:
                    decoded_prediction.append(p)
                previous = p
            hypothesis = ''.join([self.labels_map[c] for c in decoded_prediction])
            hypotheses.append(hypothesis)
        return hypotheses


def conv1d1(in_channels, out_channels, stride=1, groups=1, bias=False):
    """
    1-dim kernel version of the 1D convolution layer.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int, default 1
        Strides of the convolution.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    """
    return nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, groups=groups, bias=bias)


def conv5x5_block(in_channels, out_channels, stride=1, padding=2, dilation=1, groups=1, bias=False, use_bn=True, bn_eps=1e-05, activation=lambda : nn.ReLU(inplace=True)):
    """
    5x5 version of the standard convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int, or tuple/list of 2 int, or tuple/list of 4 int, default 2
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    """
    return ConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=5, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias, use_bn=use_bn, bn_eps=bn_eps, activation=activation)


def mixconv1x1_block(in_channels, out_channels, kernel_count, stride=1, groups=1, bias=False, use_bn=True, bn_eps=1e-05, activation=lambda : nn.ReLU(inplace=True)):
    """
    1x1 version of the mixed convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    kernel_count : int
        Kernel count.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str, or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    """
    return MixConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=[1] * kernel_count, stride=stride, padding=[0] * kernel_count, groups=groups, bias=bias, use_bn=use_bn, bn_eps=bn_eps, activation=activation)


def nasnet_avgpool1x1_s2():
    """
    NASNet specific 1x1 Average pooling layer with stride 2.
    """
    return nn.AvgPool2d(kernel_size=1, stride=2, count_include_pad=False)


def dws_branch_k3_s1_p1(in_channels, out_channels, extra_padding=False):
    """
    3x3/1/1 version of the NASNet specific depthwise separable convolution branch.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    extra_padding : bool, default False
        Whether to use extra padding.
    """
    return DwsBranch(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, extra_padding=extra_padding)


def dws_branch_k5_s2_p2(in_channels, out_channels, extra_padding=False, stem=False):
    """
    5x5/2/2 version of the NASNet specific depthwise separable convolution branch.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    extra_padding : bool, default False
        Whether to use extra padding.
    stem : bool, default False
        Whether to use squeeze reduction if False.
    """
    return DwsBranch(in_channels=in_channels, out_channels=out_channels, kernel_size=5, stride=2, padding=2, extra_padding=extra_padding, stem=stem)


def dws_branch_k7_s2_p3(in_channels, out_channels, extra_padding=False, stem=False):
    """
    7x7/2/3 version of the NASNet specific depthwise separable convolution branch.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    extra_padding : bool, default False
        Whether to use extra padding.
    stem : bool, default False
        Whether to use squeeze reduction if False.
    """
    return DwsBranch(in_channels=in_channels, out_channels=out_channels, kernel_size=7, stride=2, padding=3, extra_padding=extra_padding, stem=stem)


def nas_conv1x1(in_channels, out_channels):
    """
    1x1 version of the NASNet specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    """
    return NasConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0, groups=1)


def nasnet_avgpool3x3_s1():
    """
    NASNet specific 3x3 Average pooling layer with stride 1.
    """
    return nn.AvgPool2d(kernel_size=3, stride=1, padding=1, count_include_pad=False)


def dws_branch_k5_s1_p2(in_channels, out_channels, extra_padding=False):
    """
    5x5/1/2 version of the NASNet specific depthwise separable convolution branch.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    extra_padding : bool, default False
        Whether to use extra padding.
    """
    return DwsBranch(in_channels=in_channels, out_channels=out_channels, kernel_size=5, stride=1, padding=2, extra_padding=extra_padding)


def hard_nms(cdds, top_n=10, iou_thresh=0.25):
    """
    Hard Non-Maximum Suppression.

    Parameters:
    ----------
    cdds : np.array
        Borders.
    top_n : int, default 10
        Number of top-K informative regions.
    iou_thresh : float, default 0.25
        IoU threshold.

    Returns:
    -------
    np.array
        Filtered borders.
    """
    assert type(cdds) == np.ndarray
    assert len(cdds.shape) == 2
    assert cdds.shape[1] >= 5
    cdds = cdds.copy()
    indices = np.argsort(cdds[:, 0])
    cdds = cdds[indices]
    cdd_results = []
    res = cdds
    while res.any():
        cdd = res[-1]
        cdd_results.append(cdd)
        if len(cdd_results) == top_n:
            return np.array(cdd_results)
        res = res[:-1]
        start_max = np.maximum(res[:, 1:3], cdd[1:3])
        end_min = np.minimum(res[:, 3:5], cdd[3:5])
        lengths = end_min - start_max
        intersec_map = lengths[:, 0] * lengths[:, 1]
        intersec_map[np.logical_or(lengths[:, 0] < 0, lengths[:, 1] < 0)] = 0
        iou_map_cur = intersec_map / ((res[:, 3] - res[:, 1]) * (res[:, 4] - res[:, 2]) + (cdd[3] - cdd[1]) * (cdd[4] - cdd[2]) - intersec_map)
        res = res[iou_map_cur < iou_thresh]
    return np.array(cdd_results)


def oct_conv3x3_block(in_channels, out_channels, stride=1, padding=1, dilation=1, groups=1, bias=False, oct_alpha=0.0, oct_mode='std', bn_eps=1e-05, activation=lambda : nn.ReLU(inplace=True), activate=True):
    """
    3x3 version of the octave convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    oct_alpha : float, default 0.0
        Octave alpha coefficient.
    oct_mode : str, default 'std'
        Octave convolution mode. It can be 'first', 'norm', 'last', or 'std'.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    activate : bool, default True
        Whether activate the convolution block.
    """
    return OctConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias, oct_alpha=oct_alpha, oct_mode=oct_mode, bn_eps=bn_eps, activation=activation, activate=activate)


def oct_conv1x1_block(in_channels, out_channels, stride=1, groups=1, bias=False, oct_alpha=0.0, oct_mode='std', bn_eps=1e-05, activation=lambda : nn.ReLU(inplace=True), activate=True):
    """
    1x1 version of the octave convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    oct_alpha : float, default 0.0
        Octave alpha coefficient.
    oct_mode : str, default 'std'
        Octave convolution mode. It can be 'first', 'norm', 'last', or 'std'.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    activate : bool, default True
        Whether activate the convolution block.
    """
    return OctConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=0, groups=groups, bias=bias, oct_alpha=oct_alpha, oct_mode=oct_mode, bn_eps=bn_eps, activation=activation, activate=activate)


def poly_conv1x1(in_channels, out_channels, num_blocks):
    """
    1x1 version of the PolyNet specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    num_blocks : int
        Number of blocks (BatchNorm layers).
    """
    return PolyConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0, num_blocks=num_blocks)


def poly_res_b_block():
    """
    PolyNet type PolyResidual-Res-B block.
    """
    return conv1x1_block(in_channels=384, out_channels=1152, stride=1, activation=None)


def poly_res_c_block():
    """
    PolyNet type PolyResidual-Res-C block.
    """
    return conv1x1_block(in_channels=448, out_channels=2048, stride=1, activation=None)


def conv4x4_block(in_channels, out_channels, stride=1, padding=(1, 2, 1, 2), dilation=1, groups=1, bias=False, use_bn=True, bn_eps=1e-05, activation=lambda : nn.ReLU(inplace=True)):
    """
    4x4 version of the standard convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int, or tuple/list of 2 int, or tuple/list of 4 int, default (1, 2, 1, 2)
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    """
    return ConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=4, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias, use_bn=use_bn, bn_eps=bn_eps, activation=activation)


def deconv4x4_block(in_channels, out_channels, stride=1, padding=3, ext_padding=(2, 1, 2, 1), out_padding=0, dilation=1, groups=1, bias=False, use_bn=True, bn_eps=1e-05, activation=lambda : nn.ReLU(inplace=True)):
    """
    4x4 version of the standard deconvolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default (2, 1, 2, 1)
        Padding value for deconvolution layer.
    ext_padding : tuple/list of 4 int, default None
        Extra padding value for deconvolution layer.
    out_padding : int or tuple/list of 2 int
        Output padding value for deconvolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    use_bn : bool, default True
        Whether to use BatchNorm layer.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    """
    return DeconvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=4, stride=stride, padding=padding, ext_padding=ext_padding, out_padding=out_padding, dilation=dilation, groups=groups, bias=bias, use_bn=use_bn, bn_eps=bn_eps, activation=activation)


def saconv3x3_block(in_channels, out_channels, stride=1, padding=1, **kwargs):
    """
    3x3 version of the Split-Attention convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    """
    return SAConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=padding, **kwargs)


use_context_mans = int(torch.__version__[0]) * 100 + int(torch.__version__[2]) - (1 if 'a' in torch.__version__ else 0) > 3


class ReversibleBlockFunction(torch.autograd.Function):
    """
    RevNet reversible block function.
    """

    @staticmethod
    def forward(ctx, x, fm, gm, *params):
        with torch.no_grad():
            x1, x2 = torch.chunk(x, chunks=2, dim=1)
            x1 = x1.contiguous()
            x2 = x2.contiguous()
            y1 = x1 + fm(x2)
            y2 = x2 + gm(y1)
            y = torch.cat((y1, y2), dim=1)
            x1.set_()
            x2.set_()
            y1.set_()
            y2.set_()
            del x1, x2, y1, y2
        ctx.save_for_backward(x, y)
        ctx.fm = fm
        ctx.gm = gm
        return y

    @staticmethod
    def backward(ctx, grad_y):
        fm = ctx.fm
        gm = ctx.gm
        x, y = ctx.saved_variables
        y1, y2 = torch.chunk(y, chunks=2, dim=1)
        y1 = y1.contiguous()
        y2 = y2.contiguous()
        with torch.no_grad():
            y1_z = Variable(y1.data, requires_grad=True)
            x2 = y2 - gm(y1_z)
            x1 = y1 - fm(x2)
        with set_grad_enabled(True):
            x1_ = Variable(x1.data, requires_grad=True)
            x2_ = Variable(x2.data, requires_grad=True)
            y1_ = x1_ + fm.forward(x2_)
            y2_ = x2_ + gm(y1_)
            y = torch.cat((y1_, y2_), dim=1)
            dd = torch.autograd.grad(y, (x1_, x2_) + tuple(gm.parameters()) + tuple(fm.parameters()), grad_y)
            gm_params_len = len([p for p in gm.parameters()])
            gm_params_grads = dd[2:2 + gm_params_len]
            fm_params_grads = dd[2 + gm_params_len:]
            grad_x = torch.cat((dd[0], dd[1]), dim=1)
            y1_.detach_()
            y2_.detach_()
            del y1_, y2_
        x.data.set_(torch.cat((x1, x2), dim=1).data.contiguous())
        return (grad_x, None, None) + fm_params_grads + gm_params_grads


class ShakeDrop(torch.autograd.Function):
    """
    ShakeDrop function.
    """

    @staticmethod
    def forward(ctx, x, b, alpha):
        y = (b + alpha - b * alpha) * x
        ctx.save_for_backward(b)
        return y

    @staticmethod
    def backward(ctx, dy):
        beta = torch.rand(dy.size(0), dtype=dy.dtype, device=dy.device).view(-1, 1, 1, 1)
        b, = ctx.saved_tensors
        return (b + beta - b * beta) * dy, None, None


class ShakeShake(torch.autograd.Function):
    """
    Shake-Shake function.
    """

    @staticmethod
    def forward(ctx, x1, x2, alpha):
        y = alpha * x1 + (1 - alpha) * x2
        return y

    @staticmethod
    def backward(ctx, dy):
        beta = torch.rand(dy.size(0), dtype=dy.dtype, device=dy.device).view(-1, 1, 1, 1)
        return beta * dy, (1 - beta) * dy, None


def sha_conv3x3_block(in_channels, out_channels, stride=1, padding=1, dilation=1, groups=1, bias=False, activation=lambda : nn.ReLU(inplace=True), activate=True, shared_conv=None):
    """
    3x3 version of the shared convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    activation : function or str or None, default nn.ReLU(inplace=True)
        Activation function or name of activation function.
    activate : bool, default True
        Whether activate the convolution block.
    shared_conv : Module, default None
        Shared convolution layer.
    """
    return ShaConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias, activation=activation, activate=activate, shared_conv=shared_conv)


class SBStage(nn.Layer):
    """
    SB stage.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    down_channels : int
        Number of output channels for a downscale block.
    channels_list : list of int
        Number of output channels for all residual block.
    kernel_sizes_list : list of int
        Convolution window size for branches.
    scale_factors_list : list of int
        Scale factor for branches.
    use_residual_list : list of int
        List of flags for using residual in each ESP-block.
    se_reduction : int
        Squeeze reduction value (0 means no-se).
    in_size : tuple of 2 int
        Spatial size of the output tensor for the bilinear upsampling operation.
    bn_eps : float
        Small float added to variance in Batch norm.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, in_channels, down_channels, channels_list, kernel_sizes_list, scale_factors_list, use_residual_list, se_reduction, in_size, bn_eps, data_format='channels_last', **kwargs):
        super(SBStage, self).__init__(**kwargs)
        self.data_format = data_format
        self.down_conv = dwsconv3x3_block(in_channels=in_channels, out_channels=down_channels, strides=2, dw_use_bn=False, bn_eps=bn_eps, dw_activation=None, pw_activation=lambda : PReLU2(down_channels, data_format=data_format, name='activ'), se_reduction=se_reduction, data_format=data_format, name='down_conv')
        in_channels = down_channels
        self.main_branch = SimpleSequential(name='main_branch')
        for i, out_channels in enumerate(channels_list):
            use_residual = use_residual_list[i] == 1
            kernel_sizes = kernel_sizes_list[i]
            scale_factors = scale_factors_list[i]
            self.main_branch.add(ESPBlock(in_channels=in_channels, out_channels=out_channels, kernel_sizes=kernel_sizes, scale_factors=scale_factors, use_residual=use_residual, in_size=(in_size[0] // 2, in_size[1] // 2) if in_size else None, bn_eps=bn_eps, data_format=data_format, name='block{}'.format(i + 1)))
            in_channels = out_channels
        self.preactiv = PreActivation(in_channels=down_channels + in_channels, bn_eps=bn_eps, data_format=data_format, name='preactiv')

    def call(self, x, training=None):
        x = self.down_conv(x, training=None)
        y = self.main_branch(x, training=None)
        x = tf.concat([x, y], axis=get_channel_axis(self.data_format))
        x = self.preactiv(x, training=None)
        return x, y


class SBEncoderInitBlock(nn.Layer):
    """
    SB encoder specific initial block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    mid_channels : int
        Number of middle channels.
    out_channels : int
        Number of output channels.
    bn_eps : float, default 1e-5
        Small float added to variance in Batch norm.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, in_channels, mid_channels, out_channels, bn_eps, data_format='channels_last', **kwargs):
        super(SBEncoderInitBlock, self).__init__(**kwargs)
        self.conv1 = conv3x3_block(in_channels=in_channels, out_channels=mid_channels, strides=2, bn_eps=bn_eps, activation=lambda : PReLU2(mid_channels, data_format=data_format, name='activ'), data_format=data_format, name='conv1')
        self.conv2 = dwsconv3x3_block(in_channels=mid_channels, out_channels=out_channels, strides=2, dw_use_bn=False, bn_eps=bn_eps, dw_activation=None, pw_activation=lambda : PReLU2(out_channels, data_format=data_format, name='activ'), se_reduction=1, data_format=data_format, name='conv2')

    def call(self, x, training=None):
        x = self.conv1(x, training=None)
        x = self.conv2(x, training=None)
        return x


class SBEncoder(nn.Layer):
    """
    SB encoder for SINet.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of input channels.
    init_block_channels : list int
        Number of output channels for convolutions in the initial block.
    down_channels_list : list of int
        Number of downsample channels for each residual block.
    channels_list : list of list of int
        Number of output channels for all residual block.
    kernel_sizes_list : list of list of int
        Convolution window size for each residual block.
    scale_factors_list : list of list of int
        Scale factor for each residual block.
    use_residual_list : list of list of int
        List of flags for using residual in each residual block.
    in_size : tuple of 2 int
        Spatial size of the output tensor for the bilinear upsampling operation.
    bn_eps : float
        Small float added to variance in Batch norm.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, in_channels, out_channels, init_block_channels, down_channels_list, channels_list, kernel_sizes_list, scale_factors_list, use_residual_list, in_size, bn_eps, data_format='channels_last', **kwargs):
        super(SBEncoder, self).__init__(**kwargs)
        self.init_block = SBEncoderInitBlock(in_channels=in_channels, mid_channels=init_block_channels[0], out_channels=init_block_channels[1], bn_eps=bn_eps, data_format=data_format, name='init_block')
        in_channels = init_block_channels[1]
        self.stage1 = SBStage(in_channels=in_channels, down_channels=down_channels_list[0], channels_list=channels_list[0], kernel_sizes_list=kernel_sizes_list[0], scale_factors_list=scale_factors_list[0], use_residual_list=use_residual_list[0], se_reduction=1, in_size=(in_size[0] // 4, in_size[1] // 4) if in_size else None, bn_eps=bn_eps, data_format=data_format, name='stage1')
        in_channels = down_channels_list[0] + channels_list[0][-1]
        self.stage2 = SBStage(in_channels=in_channels, down_channels=down_channels_list[1], channels_list=channels_list[1], kernel_sizes_list=kernel_sizes_list[1], scale_factors_list=scale_factors_list[1], use_residual_list=use_residual_list[1], se_reduction=2, in_size=(in_size[0] // 8, in_size[1] // 8) if in_size else None, bn_eps=bn_eps, data_format=data_format, name='stage2')
        in_channels = down_channels_list[1] + channels_list[1][-1]
        self.output_conv = conv1x1(in_channels=in_channels, out_channels=out_channels, data_format=data_format, name='output')

    def call(self, x, training=None):
        y1 = self.init_block(x, training=None)
        x, y2 = self.stage1(y1, training=None)
        x, _ = self.stage2(x, training=None)
        x = self.output_conv(x)
        return x, y2, y1


class SBDecodeBlock(nn.Layer):
    """
    SB decoder block for SINet.

    Parameters:
    ----------
    channels : int
        Number of output classes.
    out_size : tuple of 2 int
        Spatial size of the output tensor for the bilinear upsampling operation.
    bn_eps : float
        Small float added to variance in Batch norm.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, channels, out_size, bn_eps, data_format='channels_last', **kwargs):
        super(SBDecodeBlock, self).__init__(**kwargs)
        assert channels is not None
        self.data_format = data_format
        self.up = InterpolationBlock(scale_factor=2, out_size=out_size, data_format=data_format, name='up')
        self.bn = BatchNorm(epsilon=bn_eps, data_format=data_format, name='bn')

    def call(self, x, y, training=None):
        x = self.up(x)
        x = self.bn(x, training=None)
        w_conf = tf.nn.softmax(x)
        axis = get_channel_axis(self.data_format)
        w_max = tf.broadcast_to(tf.expand_dims(tf.reduce_max(w_conf, axis=axis), axis=axis), shape=x.shape)
        x = y * (1 - w_max) + x
        return x


class SBDecoder(nn.Layer):
    """
    SB decoder for SINet.

    Parameters:
    ----------
    dim2 : int
        Size of dimension #2.
    classes : int
        Number of segmentation classes.
    out_size : tuple of 2 int
        Spatial size of the output tensor for the bilinear upsampling operation.
    bn_eps : float
        Small float added to variance in Batch norm.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, dim2, classes, out_size, bn_eps, data_format='channels_last', **kwargs):
        super(SBDecoder, self).__init__(**kwargs)
        self.decode1 = SBDecodeBlock(channels=classes, out_size=(out_size[0] // 8, out_size[1] // 8) if out_size else None, bn_eps=bn_eps, data_format=data_format, name='decode1')
        self.decode2 = SBDecodeBlock(channels=classes, out_size=(out_size[0] // 4, out_size[1] // 4) if out_size else None, bn_eps=bn_eps, data_format=data_format, name='decode2')
        self.conv3c = conv1x1_block(in_channels=dim2, out_channels=classes, bn_eps=bn_eps, activation=lambda : PReLU2(classes, data_format=data_format, name='activ'), data_format=data_format, name='conv3c')
        self.output_conv = nn.Conv2DTranspose(filters=classes, kernel_size=2, strides=2, padding='valid', output_padding=0, use_bias=False, data_format=data_format, name='output_conv')
        self.up = InterpolationBlock(scale_factor=2, out_size=out_size, data_format=data_format, name='up')

    def call(self, y3, y2, y1, training=None):
        y2 = self.conv3c(y2, training=None)
        x = self.decode1(y3, y2, training=None)
        x = self.decode2(x, y1, training=None)
        x = self.output_conv(x, training=None)
        x = self.up(x)
        return x


class SINet(tf.keras.Model):
    """
    SINet model from 'SINet: Extreme Lightweight Portrait Segmentation Networks with Spatial Squeeze Modules and
    Information Blocking Decoder,' https://arxiv.org/abs/1911.09099.

    Parameters:
    ----------
    down_channels_list : list of int
        Number of downsample channels for each residual block.
    channels_list : list of list of int
        Number of output channels for all residual block.
    kernel_sizes_list : list of list of int
        Convolution window size for each residual block.
    scale_factors_list : list of list of int
        Scale factor for each residual block.
    use_residual_list : list of list of int
        List of flags for using residual in each residual block.
    dim2 : int
        Size of dimension #2.
    bn_eps : float
        Small float added to variance in Batch norm.
    aux : bool, default False
        Whether to output an auxiliary result.
    fixed_size : bool, default True
        Whether to expect fixed spatial size of input image.
    in_channels : int, default 3
        Number of input channels.
    in_size : tuple of two ints, default (1024, 2048)
        Spatial size of the expected input image.
    classes : int, default 21
        Number of segmentation classes.
    data_format : str, default 'channels_last'
        The ordering of the dimensions in tensors.
    """

    def __init__(self, down_channels_list, channels_list, kernel_sizes_list, scale_factors_list, use_residual_list, dim2, bn_eps, aux=False, fixed_size=True, in_channels=3, in_size=(1024, 2048), classes=21, data_format='channels_last', **kwargs):
        super(SINet, self).__init__(**kwargs)
        assert fixed_size is not None
        assert in_channels > 0
        assert in_size[0] % 64 == 0 and in_size[1] % 64 == 0
        self.in_size = in_size
        self.classes = classes
        self.data_format = data_format
        self.aux = aux
        init_block_channels = [16, classes]
        out_channels = classes
        self.encoder = SBEncoder(in_channels=in_channels, out_channels=out_channels, init_block_channels=init_block_channels, down_channels_list=down_channels_list, channels_list=channels_list, kernel_sizes_list=kernel_sizes_list, scale_factors_list=scale_factors_list, use_residual_list=use_residual_list, in_size=in_size if fixed_size else None, bn_eps=bn_eps, data_format=data_format, name='encoder')
        self.decoder = SBDecoder(dim2=dim2, classes=classes, out_size=in_size if fixed_size else None, bn_eps=bn_eps, data_format=data_format, name='decoder')

    def call(self, x, training=None):
        y3, y2, y1 = self.encoder(x, training=None)
        x = self.decoder(y3, y2, y1, training=None)
        if self.aux:
            return x, y3
        else:
            return x


def sparsenet_exponential_fetch(lst):
    """
    SparseNet's specific exponential fetch.

    Parameters:
    ----------
    lst : list
        List of something.

    Returns:
    -------
    list
        Filtered list.
    """
    return [lst[len(lst) - 2 ** i] for i in range(1 + math.floor(math.log(len(lst), 2)))]


def anti_aliased_downsample(x):
    """
    Anti-Aliased Downsample operation.

    Parameters:
    ----------
    x : Tensor
        Input tensor.

    Returns:
    -------
    Tensor
        Resulted tensor.
    """
    channels = x.shape[1]
    weight = torch.tensor([1.0, 2.0, 1.0], dtype=x.dtype, device=x.device)
    weight = weight[:, None] * weight[None, :]
    weight = weight / torch.sum(weight)
    weight = weight[None, None, :, :].repeat((channels, 1, 1, 1))
    x_pad = F.pad(x, pad=(1, 1, 1, 1), mode='reflect')
    x = F.conv2d(x_pad, weight, stride=2, padding=0, groups=channels)
    return x


def space_to_depth(x):
    """
    Space-to-Depth operation.

    Parameters:
    ----------
    x : Tensor
        Input tensor.

    Returns:
    -------
    Tensor
        Resulted tensor.
    """
    k = 4
    batch, channels, height, width = x.size()
    new_height = height // k
    new_width = width // k
    new_channels = channels * k * k
    x = x.view(batch, channels, new_height, k, new_width, k)
    x = x.permute(0, 3, 5, 1, 2, 4).contiguous()
    x = x.view(batch, new_channels, new_height, new_width)
    return x


def wrn_conv1x1(in_channels, out_channels, stride, activate):
    """
    1x1 version of the WRN specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int
        Strides of the convolution.
    activate : bool
        Whether activate the convolution block.
    """
    return WRNConv(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=0, activate=activate)


def wrn_conv3x3(in_channels, out_channels, stride, activate):
    """
    3x3 version of the WRN specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int
        Strides of the convolution.
    activate : bool
        Whether activate the convolution block.
    """
    return WRNConv(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=1, activate=activate)


class Binarize(torch.autograd.Function):
    """
    Fake sign op for 1-bit weights.
    """

    @staticmethod
    def forward(ctx, x):
        return math.sqrt(2.0 / (x.shape[1] * x.shape[2] * x.shape[3])) * x.sign()

    @staticmethod
    def backward(ctx, dy):
        return dy


def pre_conv3x3_block_1bit(in_channels, out_channels, stride=1, padding=1, dilation=1, bn_affine=True, return_preact=False, activate=True, binarized=False):
    """
    3x3 version of the pre-activated convolution block with binarization.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    bn_affine : bool, default True
        Whether the BatchNorm layer learns affine parameters.
    return_preact : bool, default False
        Whether return pre-activation.
    activate : bool, default True
        Whether activate the convolution block.
    binarized : bool, default False
        Whether to use binarization.
    """
    return PreConvBlock1bit(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=padding, dilation=dilation, bn_affine=bn_affine, return_preact=return_preact, activate=activate, binarized=binarized)


def conv1x1_block_1bit(in_channels, out_channels, stride=1, padding=0, groups=1, bias=False, bn_affine=True, activate=True, binarized=False):
    """
    1x1 version of the standard convolution block with binarization.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 0
        Padding value for convolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    bn_affine : bool, default True
        Whether the BatchNorm layer learns affine parameters.
    activate : bool, default True
        Whether activate the convolution block.
    binarized : bool, default False
        Whether to use binarization.
    """
    return ConvBlock1bit(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=padding, groups=groups, bias=bias, bn_affine=bn_affine, activate=activate, binarized=binarized)


def conv3x3_1bit(in_channels, out_channels, stride=1, padding=1, dilation=1, groups=1, bias=False, binarized=False):
    """
    Convolution 3x3 layer with binarization.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    groups : int, default 1
        Number of groups.
    bias : bool, default False
        Whether the layer uses a bias vector.
    binarized : bool, default False
        Whether to use binarization.
    """
    return Conv2d1bit(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias, binarized=binarized)


def dws_conv3x3_block(in_channels, out_channels, activate):
    """
    3x3 version of the depthwise separable convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    activate : bool
        Whether activate the convolution block.
    """
    return DwsConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, activate=activate)


def pre_xconv1x1_block(in_channels, out_channels, stride=1, bias=False, return_preact=False, activate=True, expand_ratio=2):
    """
    1x1 version of the pre-activated x-convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    bias : bool, default False
        Whether the layer uses a bias vector.
    return_preact : bool, default False
        Whether return pre-activation.
    activate : bool, default True
        Whether activate the convolution block.
    expand_ratio : int, default 2
        Ratio of expansion.
    """
    return PreXConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=0, bias=bias, return_preact=return_preact, activate=activate, expand_ratio=expand_ratio)


def pre_xconv3x3_block(in_channels, out_channels, stride=1, padding=1, dilation=1, return_preact=False, activate=True, expand_ratio=2):
    """
    3x3 version of the pre-activated x-convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    stride : int or tuple/list of 2 int, default 1
        Strides of the convolution.
    padding : int or tuple/list of 2 int, default 1
        Padding value for convolution layer.
    dilation : int or tuple/list of 2 int, default 1
        Dilation value for convolution layer.
    return_preact : bool, default False
        Whether return pre-activation.
    activate : bool, default True
        Whether activate the convolution block.
    expand_ratio : int, default 2
        Ratio of expansion.
    """
    return PreXConvBlock(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=padding, dilation=dilation, return_preact=return_preact, activate=activate, expand_ratio=expand_ratio)


class PytorchModel(torch.nn.Module):

    def __init__(self):
        super(PytorchModel, self).__init__()
        self.dense = torch.nn.Linear(in_features=1024, out_features=1000, bias=False)

    def forward(self, x):
        x = self.dense(x)
        return x

