[
  {
    "path": "README.md",
    "content": "# efficientdet\nBiFPN and Modified BiFPN.\n\neffcientNet backbones and pretrained weights from @rwightman(https://github.com/rwightman/gen-efficientnet-pytorch)\n\n# TODO\ntrain and test\n"
  },
  {
    "path": "backbones/efficientnet.py",
    "content": "import torch.nn as nn\n\nfrom torch.nn.modules.batchnorm import _BatchNorm\nfrom ..registry import BACKBONES\nimport sys\nsys.path.append('./mmdet/models/backbones')\nimport geffnet\n\n@BACKBONES.register_module\nclass EfficientNet(nn.Module):\n    \"\"\"EfficientNet backbone and pretrained from https://github.com/rwightman/gen-efficientnet-pytorch\n\n    Args:\n        model_name (string): tf_efficientnet_b0-b7.\n        pretrained (bool) : load pretrained weights, must be True.\n        out_indices (Sequence[int]): Output from which stages. Should be (2, 3, 4, 5, 6) in EfficientDet.\n        style (str): `pytorch` or `caffe`. If set to \"pytorch\", the stride-two\n            layer is the 3x3 conv layer, otherwise the stride-two layer is\n            the first 1x1 conv layer. Not used.\n        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).\n            -1 means not freezing any parameters.\n\n    Example:\n        >>> from mmdet.models import EfficientNet\n        >>> import torch\n        >>> self = EfficientNet(model_name='tf_efficientnet_b2', pretrained=False)\n        >>> self.eval()\n        >>> inputs = torch.rand(1,3,768,768)\n        >>> level_outputs = self(inputs)\n        >>> for level_out in level_outputs:\n        ...     print(tuple(level_out.shape))\n        (1, 48, 96, 96)\n        (1, 88, 48, 48)\n        (1, 120, 24, 24)\n        (1, 208, 12, 12)\n        (1, 352, 6, 6)\n    \"\"\"\n\n    def __init__(self,\n                 model_name,\n                 pretrained=True,\n                 out_indices=(2, 3, 4, 5, 6),\n                 style='pytorch',\n                 frozen_stages=-1,\n                 norm_eval=True):\n        super(EfficientNet, self).__init__()\n        self.out_indices = out_indices\n        self.style = style\n        self.frozen_stages = frozen_stages\n        self.norm_eval = norm_eval\n        self.model = geffnet.create_model(model_name,pretrained=pretrained)\n\n        self._freeze_stages()\n\n    def _freeze_stages(self):\n        if self.frozen_stages >= 0:\n            self.eval()\n            for param in self.parameters():\n                param.requires_grad = False\n\n    def init_weights(self, pretrained=None):\n        return\n\n    def forward(self, x):\n        feature_map = self.model(x)\n        outs=[]\n        for i in self.out_indices:\n            outs.append(feature_map[i])\n        return tuple(outs)\n\n    def train(self, mode=True): #need modify\n        super(EfficientNet, self).train(mode)\n        self._freeze_stages()\n        if mode and self.norm_eval:\n            for m in self.modules():\n                # trick: eval have effect on BatchNorm only\n                if isinstance(m, _BatchNorm):\n                    m.eval()\n"
  },
  {
    "path": "backbones/geffnet/__init__.py",
    "content": "from .gen_efficientnet import *\nfrom .mobilenetv3 import *\nfrom .model_factory import create_model\nfrom .config import is_exportable, is_scriptable, set_exportable, set_scriptable\nfrom .activations import *"
  },
  {
    "path": "backbones/geffnet/activations/__init__.py",
    "content": "from geffnet import config\nfrom geffnet.activations.activations_autofn import *\nfrom geffnet.activations.activations_jit import *\nfrom geffnet.activations.activations import *\n\n\n_ACT_FN_DEFAULT = dict(\n    swish=swish,\n    mish=mish,\n    relu=F.relu,\n    relu6=F.relu6,\n    sigmoid=sigmoid,\n    tanh=tanh,\n    hard_sigmoid=hard_sigmoid,\n    hard_swish=hard_swish,\n)\n\n_ACT_FN_AUTO = dict(\n    swish=swish_auto,\n    mish=mish_auto,\n)\n\n_ACT_FN_JIT = dict(\n    swish=swish_jit,\n    mish=mish_jit,\n    #hard_swish=hard_swish_jit,\n    #hard_sigmoid_jit=hard_sigmoid_jit,\n)\n\n_ACT_LAYER_DEFAULT = dict(\n    swish=Swish,\n    mish=Mish,\n    relu=nn.ReLU,\n    relu6=nn.ReLU6,\n    sigmoid=Sigmoid,\n    tanh=Tanh,\n    hard_sigmoid=HardSigmoid,\n    hard_swish=HardSwish,\n)\n\n_ACT_LAYER_AUTO = dict(\n    swish=SwishAuto,\n    mish=MishAuto,\n)\n\n_ACT_LAYER_JIT = dict(\n    swish=SwishJit,\n    mish=MishJit,\n    #hard_swish=HardSwishJit,\n    #hard_sigmoid=HardSigmoidJit\n)\n\n_OVERRIDE_FN = dict()\n_OVERRIDE_LAYER = dict()\n\n\ndef add_override_act_fn(name, fn):\n    global _OVERRIDE_FN\n    _OVERRIDE_FN[name] = fn\n\n\ndef update_override_act_fn(overrides):\n    assert isinstance(overrides, dict)\n    global _OVERRIDE_FN\n    _OVERRIDE_FN.update(overrides)\n\n\ndef clear_override_act_fn():\n    global _OVERRIDE_FN\n    _OVERRIDE_FN = dict()\n\n\ndef add_override_act_layer(name, fn):\n    _OVERRIDE_LAYER[name] = fn\n\n\ndef update_override_act_layer(overrides):\n    assert isinstance(overrides, dict)\n    global _OVERRIDE_LAYER\n    _OVERRIDE_LAYER.update(overrides)\n\n\ndef clear_override_act_layer():\n    global _OVERRIDE_LAYER\n    _OVERRIDE_LAYER = dict()\n\n\ndef get_act_fn(name='relu'):\n    \"\"\" Activation Function Factory\n    Fetching activation fns by name with this function allows export or torch script friendly\n    functions to be returned dynamically based on current config.\n    \"\"\"\n    if name in _OVERRIDE_FN:\n        return _OVERRIDE_FN[name]\n    if not config.is_exportable() and not config.is_scriptable():\n        # If not exporting or scripting the model, first look for a JIT optimized version\n        # of our activation, then a custom autograd.Function variant before defaulting to\n        # a Python or Torch builtin impl\n        if name in _ACT_FN_JIT:\n            return _ACT_FN_JIT[name]\n        if name in _ACT_FN_AUTO:\n            return _ACT_FN_AUTO[name]\n    return _ACT_FN_DEFAULT[name]\n\n\ndef get_act_layer(name='relu'):\n    \"\"\" Activation Layer Factory\n    Fetching activation layers by name with this function allows export or torch script friendly\n    functions to be returned dynamically based on current config.\n    \"\"\"\n    if name in _OVERRIDE_LAYER:\n        return _OVERRIDE_LAYER[name]\n    if not config.is_exportable() and not config.is_scriptable():\n        if name in _ACT_LAYER_JIT:\n            return _ACT_LAYER_JIT[name]\n        if name in _ACT_LAYER_AUTO:\n            return _ACT_LAYER_AUTO[name]\n    return _ACT_LAYER_DEFAULT[name]\n\n\n"
  },
  {
    "path": "backbones/geffnet/activations/activations.py",
    "content": "from torch import nn as nn\nfrom torch.nn import functional as F\n\n\ndef swish(x, inplace: bool = False):\n    \"\"\"Swish - Described in: https://arxiv.org/abs/1710.05941\n    \"\"\"\n    return x.mul_(x.sigmoid()) if inplace else x.mul(x.sigmoid())\n\n\nclass Swish(nn.Module):\n    def __init__(self, inplace: bool = False):\n        super(Swish, self).__init__()\n        self.inplace = inplace\n\n    def forward(self, x):\n        return swish(x, self.inplace)\n\n\ndef mish(x, inplace: bool = False):\n    \"\"\"Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681\n    \"\"\"\n    return x.mul(F.softplus(x).tanh())\n\n\nclass Mish(nn.Module):\n    def __init__(self, inplace: bool = False):\n        super(Mish, self).__init__()\n        self.inplace = inplace\n\n    def forward(self, x):\n        return mish(x, self.inplace)\n\n\ndef sigmoid(x, inplace: bool = False):\n    return x.sigmoid_() if inplace else x.sigmoid()\n\n\n# PyTorch has this, but not with a consistent inplace argmument interface\nclass Sigmoid(nn.Module):\n    def __init__(self, inplace: bool = False):\n        super(Sigmoid, self).__init__()\n        self.inplace = inplace\n\n    def forward(self, x):\n        return x.sigmoid_() if self.inplace else x.sigmoid()\n\n\ndef tanh(x, inplace: bool = False):\n    return x.tanh_() if inplace else x.tanh()\n\n\n# PyTorch has this, but not with a consistent inplace argmument interface\nclass Tanh(nn.Module):\n    def __init__(self, inplace: bool = False):\n        super(Tanh, self).__init__()\n        self.inplace = inplace\n\n    def forward(self, x):\n        return x.tanh_() if self.inplace else x.tanh()\n\n\ndef hard_swish(x, inplace: bool = False):\n    inner = F.relu6(x + 3.).div_(6.)\n    return x.mul_(inner) if inplace else x.mul(inner)\n\n\nclass HardSwish(nn.Module):\n    def __init__(self, inplace: bool = False):\n        super(HardSwish, self).__init__()\n        self.inplace = inplace\n\n    def forward(self, x):\n        return hard_swish(x, self.inplace)\n\n\ndef hard_sigmoid(x, inplace: bool = False):\n    if inplace:\n        return x.add_(3.).clamp_(0., 6.).div_(6.)\n    else:\n        return F.relu6(x + 3.) / 6.\n\n\nclass HardSigmoid(nn.Module):\n    def __init__(self, inplace: bool = False):\n        super(HardSigmoid, self).__init__()\n        self.inplace = inplace\n\n    def forward(self, x):\n        return hard_sigmoid(x, self.inplace)\n\n\n"
  },
  {
    "path": "backbones/geffnet/activations/activations_autofn.py",
    "content": "import torch\nfrom torch import nn as nn\nfrom torch.nn import functional as F\n\n\n__all__ = ['swish_auto', 'SwishAuto', 'mish_auto', 'MishAuto']\n\n\nclass SwishAutoFn(torch.autograd.Function):\n    \"\"\"Swish - Described in: https://arxiv.org/abs/1710.05941\n    Memory efficient variant from:\n     https://medium.com/the-artificial-impostor/more-memory-efficient-swish-activation-function-e07c22c12a76\n    \"\"\"\n    @staticmethod\n    def forward(ctx, x):\n        result = x.mul(torch.sigmoid(x))\n        ctx.save_for_backward(x)\n        return result\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        x = ctx.saved_tensors[0]\n        x_sigmoid = torch.sigmoid(x)\n        return grad_output.mul(x_sigmoid * (1 + x * (1 - x_sigmoid)))\n\n\ndef swish_auto(x, inplace=False):\n    # inplace ignored\n    return SwishAutoFn.apply(x)\n\n\nclass SwishAuto(nn.Module):\n    def __init__(self, inplace: bool = False):\n        super(SwishAuto, self).__init__()\n        self.inplace = inplace\n\n    def forward(self, x):\n        return SwishAutoFn.apply(x)\n\n\nclass MishAutoFn(torch.autograd.Function):\n    \"\"\"Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681\n    Experimental memory-efficient variant\n    \"\"\"\n\n    @staticmethod\n    def forward(ctx, x):\n        ctx.save_for_backward(x)\n        y = x.mul(torch.tanh(F.softplus(x)))  # x * tanh(ln(1 + exp(x)))\n        return y\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        x = ctx.saved_tensors[0]\n        x_sigmoid = torch.sigmoid(x)\n        x_tanh_sp = F.softplus(x).tanh()\n        return grad_output.mul(x_tanh_sp + x * x_sigmoid * (1 - x_tanh_sp * x_tanh_sp))\n\n\ndef mish_auto(x, inplace=False):\n    # inplace ignored\n    return MishAutoFn.apply(x)\n\n\nclass MishAuto(nn.Module):\n    def __init__(self, inplace: bool = False):\n        super(MishAuto, self).__init__()\n        self.inplace = inplace\n\n    def forward(self, x):\n        return MishAutoFn.apply(x)\n\n"
  },
  {
    "path": "backbones/geffnet/activations/activations_jit.py",
    "content": "import torch\nfrom torch import nn as nn\nfrom torch.nn import functional as F\n\n\n__all__ = ['swish_jit', 'SwishJit', 'mish_jit', 'MishJit']\n           #'hard_swish_jit', 'HardSwishJit', 'hard_sigmoid_jit', 'HardSigmoidJit']\n\n\n@torch.jit.script\ndef swish_jit_fwd(x):\n    return x.mul(torch.sigmoid(x))\n\n\n@torch.jit.script\ndef swish_jit_bwd(x, grad_output):\n    x_sigmoid = torch.sigmoid(x)\n    return grad_output * (x_sigmoid * (1 + x * (1 - x_sigmoid)))\n\n\nclass SwishJitAutoFn(torch.autograd.Function):\n    \"\"\" torch.jit.script optimised Swish\n    Inspired by conversation btw Jeremy Howard & Adam Pazske\n    https://twitter.com/jeremyphoward/status/1188251041835315200\n    \"\"\"\n    @staticmethod\n    def forward(ctx, x):\n        ctx.save_for_backward(x)\n        return swish_jit_fwd(x)\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        x = ctx.saved_tensors[0]\n        return swish_jit_bwd(x, grad_output)\n\n\ndef swish_jit(x, inplace=False):\n    # inplace ignored\n    return SwishJitAutoFn.apply(x)\n\n\nclass SwishJit(nn.Module):\n    def __init__(self, inplace: bool = False):\n        super(SwishJit, self).__init__()\n        self.inplace = inplace\n\n    def forward(self, x):\n        return SwishJitAutoFn.apply(x)\n\n\n@torch.jit.script\ndef mish_jit_fwd(x):\n    return x.mul(torch.tanh(F.softplus(x)))\n\n\n@torch.jit.script\ndef mish_jit_bwd(x, grad_output):\n    x_sigmoid = torch.sigmoid(x)\n    x_tanh_sp = F.softplus(x).tanh()\n    return grad_output.mul(x_tanh_sp + x * x_sigmoid * (1 - x_tanh_sp * x_tanh_sp))\n\n\nclass MishJitAutoFn(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, x):\n        ctx.save_for_backward(x)\n        return mish_jit_fwd(x)\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        x = ctx.saved_tensors[0]\n        return mish_jit_bwd(x, grad_output)\n\n\ndef mish_jit(x, inplace=False):\n    # inplace ignored\n    return MishJitAutoFn.apply(x)\n\n\nclass MishJit(nn.Module):\n    def __init__(self, inplace: bool = False):\n        super(MishJit, self).__init__()\n        self.inplace = inplace\n\n    def forward(self, x):\n        return MishJitAutoFn.apply(x)\n\n\n# @torch.jit.script\n# def hard_swish_jit(x, inplac: bool = False):\n#     return x.mul(F.relu6(x + 3.).mul_(1./6.))\n#\n#\n# class HardSwishJit(nn.Module):\n#     def __init__(self, inplace: bool = False):\n#         super(HardSwishJit, self).__init__()\n#\n#     def forward(self, x):\n#         return hard_swish_jit(x)\n#\n#\n# @torch.jit.script\n# def hard_sigmoid_jit(x, inplace: bool = False):\n#     return F.relu6(x + 3.).mul(1./6.)\n#\n#\n# class HardSigmoidJit(nn.Module):\n#     def __init__(self, inplace: bool = False):\n#         super(HardSigmoidJit, self).__init__()\n#\n#     def forward(self, x):\n#         return hard_sigmoid_jit(x)\n"
  },
  {
    "path": "backbones/geffnet/config.py",
    "content": "\"\"\" Global Config and Constants\n\"\"\"\n\n__all__ = ['is_exportable', 'is_scriptable', 'set_exportable', 'set_scriptable']\n\n# Set to True if exporting a model with Same padding via ONNX\n_EXPORTABLE = False\n\n# Set to True if wanting to use torch.jit.script on a model\n_SCRIPTABLE = False\n\n\ndef is_exportable():\n    return _EXPORTABLE\n\n\ndef set_exportable(value):\n    global _EXPORTABLE\n    _EXPORTABLE = value\n\n\ndef is_scriptable():\n    return _SCRIPTABLE\n\n\ndef set_scriptable(value):\n    global _SCRIPTABLE\n    _SCRIPTABLE = value\n\n"
  },
  {
    "path": "backbones/geffnet/conv2d_layers.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch._six import container_abcs\n\nfrom itertools import repeat\nfrom functools import partial\nfrom typing import Union, List, Tuple, Optional, Callable\nimport numpy as np\nimport math\n\nfrom .config import *\n\n\ndef _ntuple(n):\n    def parse(x):\n        if isinstance(x, container_abcs.Iterable):\n            return x\n        return tuple(repeat(x, n))\n    return parse\n\n\n_single = _ntuple(1)\n_pair = _ntuple(2)\n_triple = _ntuple(3)\n_quadruple = _ntuple(4)\n\n\ndef _is_static_pad(kernel_size, stride=1, dilation=1, **_):\n    return stride == 1 and (dilation * (kernel_size - 1)) % 2 == 0\n\n\ndef _get_padding(kernel_size, stride=1, dilation=1, **_):\n    padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2\n    return padding\n\n\ndef _calc_same_pad(i: int, k: int, s: int, d: int):\n    return max((math.ceil(i / s) - 1) * s + (k - 1) * d + 1 - i, 0)\n\n\ndef _same_pad_arg(input_size, kernel_size, stride, dilation):\n    ih, iw = input_size\n    kh, kw = kernel_size\n    pad_h = _calc_same_pad(ih, kh, stride[0], dilation[0])\n    pad_w = _calc_same_pad(iw, kw, stride[1], dilation[1])\n    return [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]\n\n\ndef _split_channels(num_chan, num_groups):\n    split = [num_chan // num_groups for _ in range(num_groups)]\n    split[0] += num_chan - sum(split)\n    return split\n\n\ndef conv2d_same(\n        x, weight: torch.Tensor, bias: Optional[torch.Tensor] = None, stride: Tuple[int, int] = (1, 1),\n        padding: Tuple[int, int] = (0, 0), dilation: Tuple[int, int] = (1, 1), groups: int = 1):\n    ih, iw = x.size()[-2:]\n    kh, kw = weight.size()[-2:]\n    pad_h = _calc_same_pad(ih, kh, stride[0], dilation[0])\n    pad_w = _calc_same_pad(iw, kw, stride[1], dilation[1])\n    if pad_h > 0 or pad_w > 0:\n        x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2])\n    return F.conv2d(x, weight, bias, stride, (0, 0), dilation, groups)\n\n\nclass Conv2dSame(nn.Conv2d):\n    \"\"\" Tensorflow like 'SAME' convolution wrapper for 2D convolutions\n    \"\"\"\n\n    # pylint: disable=unused-argument\n    def __init__(self, in_channels, out_channels, kernel_size, stride=1,\n                 padding=0, dilation=1, groups=1, bias=True):\n        super(Conv2dSame, self).__init__(\n            in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias)\n\n    def forward(self, x):\n        return conv2d_same(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)\n\n\nclass Conv2dSameExport(nn.Conv2d):\n    \"\"\" ONNX export friendly Tensorflow like 'SAME' convolution wrapper for 2D convolutions\n\n    NOTE: This does not currently work with torch.jit.script\n    \"\"\"\n\n    # pylint: disable=unused-argument\n    def __init__(self, in_channels, out_channels, kernel_size, stride=1,\n                 padding=0, dilation=1, groups=1, bias=True):\n        super(Conv2dSameExport, self).__init__(\n            in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias)\n        self.pad = None\n        self.pad_input_size = (0, 0)\n\n    def forward(self, x):\n        input_size = x.size()[-2:]\n        if self.pad is None:\n            pad_arg = _same_pad_arg(input_size, self.weight.size()[-2:], self.stride, self.dilation)\n            self.pad = nn.ZeroPad2d(pad_arg)\n            self.pad_input_size = input_size\n        else:\n            assert self.pad_input_size == input_size\n\n        x = self.pad(x)\n        return F.conv2d(\n            x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)\n\n\ndef get_padding_value(padding, kernel_size, **kwargs):\n    dynamic = False\n    if isinstance(padding, str):\n        # for any string padding, the padding will be calculated for you, one of three ways\n        padding = padding.lower()\n        if padding == 'same':\n            # TF compatible 'SAME' padding, has a performance and GPU memory allocation impact\n            if _is_static_pad(kernel_size, **kwargs):\n                # static case, no extra overhead\n                padding = _get_padding(kernel_size, **kwargs)\n            else:\n                # dynamic padding\n                padding = 0\n                dynamic = True\n        elif padding == 'valid':\n            # 'VALID' padding, same as padding=0\n            padding = 0\n        else:\n            # Default to PyTorch style 'same'-ish symmetric padding\n            padding = _get_padding(kernel_size, **kwargs)\n    return padding, dynamic\n\n\ndef create_conv2d_pad(in_chs, out_chs, kernel_size, **kwargs):\n    padding = kwargs.pop('padding', '')\n    kwargs.setdefault('bias', False)\n    padding, is_dynamic = get_padding_value(padding, kernel_size, **kwargs)\n    if is_dynamic:\n        if is_exportable():\n            assert not is_scriptable()\n            return Conv2dSameExport(in_chs, out_chs, kernel_size, **kwargs)\n        else:\n            return Conv2dSame(in_chs, out_chs, kernel_size, **kwargs)\n    else:\n        return nn.Conv2d(in_chs, out_chs, kernel_size, padding=padding, **kwargs)\n\n\nclass MixedConv2d(nn.Module):\n    \"\"\" Mixed Grouped Convolution\n    Based on MDConv and GroupedConv in MixNet impl:\n      https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mixnet/custom_layers.py\n\n    NOTE: This does not currently work with torch.jit.script\n    \"\"\"\n\n    def __init__(self, in_channels, out_channels, kernel_size=3,\n                 stride=1, padding='', dilation=1, depthwise=False, **kwargs):\n        super(MixedConv2d, self).__init__()\n\n        kernel_size = kernel_size if isinstance(kernel_size, list) else [kernel_size]\n        num_groups = len(kernel_size)\n        in_splits = _split_channels(in_channels, num_groups)\n        out_splits = _split_channels(out_channels, num_groups)\n        for idx, (k, in_ch, out_ch) in enumerate(zip(kernel_size, in_splits, out_splits)):\n            conv_groups = out_ch if depthwise else 1\n            # use add_module to keep key space clean\n            self.add_module(\n                str(idx),\n                create_conv2d_pad(\n                    in_ch, out_ch, k, stride=stride,\n                    padding=padding, dilation=dilation, groups=conv_groups, **kwargs)\n            )\n        self.splits = in_splits\n\n    def forward(self, x):\n        x_split = torch.split(x, self.splits, 1)\n        x_out = [c(x) for x, c in zip(x_split, self._modules.values())]\n        x = torch.cat(x_out, 1)\n        return x\n\n\ndef get_condconv_initializer(initializer, num_experts, expert_shape):\n    def condconv_initializer(weight):\n        \"\"\"CondConv initializer function.\"\"\"\n        num_params = np.prod(expert_shape)\n        if (len(weight.shape) != 2 or weight.shape[0] != num_experts or\n                weight.shape[1] != num_params):\n            raise (ValueError(\n                'CondConv variables must have shape [num_experts, num_params]'))\n        for i in range(num_experts):\n            initializer(weight[i].view(expert_shape))\n    return condconv_initializer\n\n\nclass CondConv2d(nn.Module):\n    \"\"\" Conditional Convolution\n    Inspired by: https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/condconv/condconv_layers.py\n\n    Grouped convolution hackery for parallel execution of the per-sample kernel filters inspired by this discussion:\n    https://github.com/pytorch/pytorch/issues/17983\n    \"\"\"\n    __constants__ = ['bias', 'in_channels', 'out_channels', 'dynamic_padding']\n\n    def __init__(self, in_channels, out_channels, kernel_size=3,\n                 stride=1, padding='', dilation=1, groups=1, bias=False, num_experts=4):\n        super(CondConv2d, self).__init__()\n\n        self.in_channels = in_channels\n        self.out_channels = out_channels\n        self.kernel_size = _pair(kernel_size)\n        self.stride = _pair(stride)\n        padding_val, is_padding_dynamic = get_padding_value(\n            padding, kernel_size, stride=stride, dilation=dilation)\n        self.dynamic_padding = is_padding_dynamic  # if in forward to work with torchscript\n        self.padding = _pair(padding_val)\n        self.dilation = _pair(dilation)\n        self.groups = groups\n        self.num_experts = num_experts\n\n        self.weight_shape = (self.out_channels, self.in_channels // self.groups) + self.kernel_size\n        weight_num_param = 1\n        for wd in self.weight_shape:\n            weight_num_param *= wd\n        self.weight = torch.nn.Parameter(torch.Tensor(self.num_experts, weight_num_param))\n\n        if bias:\n            self.bias_shape = (self.out_channels,)\n            self.bias = torch.nn.Parameter(torch.Tensor(self.num_experts, self.out_channels))\n        else:\n            self.register_parameter('bias', None)\n\n        self.reset_parameters()\n\n    def reset_parameters(self):\n        init_weight = get_condconv_initializer(\n            partial(nn.init.kaiming_uniform_, a=math.sqrt(5)), self.num_experts, self.weight_shape)\n        init_weight(self.weight)\n        if self.bias is not None:\n            fan_in = np.prod(self.weight_shape[1:])\n            bound = 1 / math.sqrt(fan_in)\n            init_bias = get_condconv_initializer(\n                partial(nn.init.uniform_, a=-bound, b=bound), self.num_experts, self.bias_shape)\n            init_bias(self.bias)\n\n    def forward(self, x, routing_weights):\n        B, C, H, W = x.shape\n        weight = torch.matmul(routing_weights, self.weight)\n        new_weight_shape = (B * self.out_channels, self.in_channels // self.groups) + self.kernel_size\n        weight = weight.view(new_weight_shape)\n        bias = None\n        if self.bias is not None:\n            bias = torch.matmul(routing_weights, self.bias)\n            bias = bias.view(B * self.out_channels)\n        # move batch elements with channels so each batch element can be efficiently convolved with separate kernel\n        x = x.view(1, B * C, H, W)\n        if self.dynamic_padding:\n            out = conv2d_same(\n                x, weight, bias, stride=self.stride, padding=self.padding,\n                dilation=self.dilation, groups=self.groups * B)\n        else:\n            out = F.conv2d(\n                x, weight, bias, stride=self.stride, padding=self.padding,\n                dilation=self.dilation, groups=self.groups * B)\n        out = out.permute([1, 0, 2, 3]).view(B, self.out_channels, out.shape[-2], out.shape[-1])\n\n        # Literal port (from TF definition)\n        # x = torch.split(x, 1, 0)\n        # weight = torch.split(weight, 1, 0)\n        # if self.bias is not None:\n        #     bias = torch.matmul(routing_weights, self.bias)\n        #     bias = torch.split(bias, 1, 0)\n        # else:\n        #     bias = [None] * B\n        # out = []\n        # for xi, wi, bi in zip(x, weight, bias):\n        #     wi = wi.view(*self.weight_shape)\n        #     if bi is not None:\n        #         bi = bi.view(*self.bias_shape)\n        #     out.append(self.conv_fn(\n        #         xi, wi, bi, stride=self.stride, padding=self.padding,\n        #         dilation=self.dilation, groups=self.groups))\n        # out = torch.cat(out, 0)\n        return out\n\n\ndef select_conv2d(in_chs, out_chs, kernel_size, **kwargs):\n    assert 'groups' not in kwargs  # only use 'depthwise' bool arg\n    if isinstance(kernel_size, list):\n        assert 'num_experts' not in kwargs  # MixNet + CondConv combo not supported currently\n        # We're going to use only lists for defining the MixedConv2d kernel groups,\n        # ints, tuples, other iterables will continue to pass to normal conv and specify h, w.\n        m = MixedConv2d(in_chs, out_chs, kernel_size, **kwargs)\n    else:\n        depthwise = kwargs.pop('depthwise', False)\n        groups = out_chs if depthwise else 1\n        if 'num_experts' in kwargs and kwargs['num_experts'] > 0:\n            m = CondConv2d(in_chs, out_chs, kernel_size, groups=groups, **kwargs)\n        else:\n            m = create_conv2d_pad(in_chs, out_chs, kernel_size, groups=groups, **kwargs)\n    return m\n"
  },
  {
    "path": "backbones/geffnet/efficientnet_builder.py",
    "content": "import re\nfrom copy import deepcopy\n\nfrom .conv2d_layers import *\nfrom geffnet.activations import *\n\n\n# Defaults used for Google/Tensorflow training of mobile networks /w RMSprop as per\n# papers and TF reference implementations. PT momentum equiv for TF decay is (1 - TF decay)\n# NOTE: momentum varies btw .99 and .9997 depending on source\n# .99 in official TF TPU impl\n# .9997 (/w .999 in search space) for paper\n#\n# PyTorch defaults are momentum = .1, eps = 1e-5\n#\nBN_MOMENTUM_TF_DEFAULT = 1 - 0.99\nBN_EPS_TF_DEFAULT = 1e-3\n_BN_ARGS_TF = dict(momentum=BN_MOMENTUM_TF_DEFAULT, eps=BN_EPS_TF_DEFAULT)\n\n\ndef get_bn_args_tf():\n    return _BN_ARGS_TF.copy()\n\n\ndef resolve_bn_args(kwargs):\n    bn_args = get_bn_args_tf() if kwargs.pop('bn_tf', False) else {}\n    bn_momentum = kwargs.pop('bn_momentum', None)\n    if bn_momentum is not None:\n        bn_args['momentum'] = bn_momentum\n    bn_eps = kwargs.pop('bn_eps', None)\n    if bn_eps is not None:\n        bn_args['eps'] = bn_eps\n    return bn_args\n\n\n_SE_ARGS_DEFAULT = dict(\n    gate_fn=sigmoid,\n    act_layer=None,  # None == use containing block's activation layer\n    reduce_mid=False,\n    divisor=1)\n\n\ndef resolve_se_args(kwargs, in_chs, act_layer=None):\n    se_kwargs = kwargs.copy() if kwargs is not None else {}\n    # fill in args that aren't specified with the defaults\n    for k, v in _SE_ARGS_DEFAULT.items():\n        se_kwargs.setdefault(k, v)\n    # some models, like MobilNetV3, calculate SE reduction chs from the containing block's mid_ch instead of in_ch\n    if not se_kwargs.pop('reduce_mid'):\n        se_kwargs['reduced_base_chs'] = in_chs\n    # act_layer override, if it remains None, the containing block's act_layer will be used\n    if se_kwargs['act_layer'] is None:\n        assert act_layer is not None\n        se_kwargs['act_layer'] = act_layer\n    return se_kwargs\n\n\ndef resolve_act_layer(kwargs, default='relu'):\n    act_layer = kwargs.pop('act_layer', default)\n    if isinstance(act_layer, str):\n        act_layer = get_act_layer(act_layer)\n    return act_layer\n\n\ndef make_divisible(v: int, divisor: int = 8, min_value: int = None):\n    min_value = min_value or divisor\n    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)\n    if new_v < 0.9 * v:  # ensure round down does not go down by more than 10%.\n        new_v += divisor\n    return new_v\n\n\ndef round_channels(channels, multiplier=1.0, divisor=8, channel_min=None):\n    \"\"\"Round number of filters based on depth multiplier.\"\"\"\n    if not multiplier:\n        return channels\n    channels *= multiplier\n    return make_divisible(channels, divisor, channel_min)\n\n\ndef drop_connect(inputs, training: bool = False, drop_connect_rate: float = 0.):\n    \"\"\"Apply drop connect.\"\"\"\n    if not training:\n        return inputs\n\n    keep_prob = 1 - drop_connect_rate\n    random_tensor = keep_prob + torch.rand(\n        (inputs.size()[0], 1, 1, 1), dtype=inputs.dtype, device=inputs.device)\n    random_tensor.floor_()  # binarize\n    output = inputs.div(keep_prob) * random_tensor\n    return output\n\n\nclass SqueezeExcite(nn.Module):\n    __constants__ = ['gate_fn']\n\n    def __init__(self, in_chs, se_ratio=0.25, reduced_base_chs=None, act_layer=nn.ReLU, gate_fn=sigmoid, divisor=1):\n        super(SqueezeExcite, self).__init__()\n        self.gate_fn = gate_fn\n        reduced_chs = make_divisible((reduced_base_chs or in_chs) * se_ratio, divisor)\n        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n        self.conv_reduce = nn.Conv2d(in_chs, reduced_chs, 1, bias=True)\n        self.act1 = act_layer(inplace=True)\n        self.conv_expand = nn.Conv2d(reduced_chs, in_chs, 1, bias=True)\n\n    def forward(self, x):\n        # tensor.view + mean bad for ONNX export (produces mess of gather ops that break TensorRT)\n        x_se = self.avg_pool(x)\n        x_se = self.conv_reduce(x_se)\n        x_se = self.act1(x_se)\n        x_se = self.conv_expand(x_se)\n        x = x * self.gate_fn(x_se)\n        return x\n\n\nclass ConvBnAct(nn.Module):\n    def __init__(self, in_chs, out_chs, kernel_size,\n                 stride=1, pad_type='', act_layer=nn.ReLU, norm_layer=nn.BatchNorm2d, norm_kwargs=None):\n        super(ConvBnAct, self).__init__()\n        assert stride in [1, 2]\n        norm_kwargs = norm_kwargs or {}\n        self.conv = select_conv2d(in_chs, out_chs, kernel_size, stride=stride, padding=pad_type)\n        self.bn1 = norm_layer(out_chs, **norm_kwargs)\n        self.act1 = act_layer(inplace=True)\n\n    def forward(self, x):\n        x = self.conv(x)\n        x = self.bn1(x)\n        x = self.act1(x)\n        return x\n\n\nclass DepthwiseSeparableConv(nn.Module):\n    \"\"\" DepthwiseSeparable block\n    Used for DS convs in MobileNet-V1 and in the place of IR blocks with an expansion\n    factor of 1.0. This is an alternative to having a IR with optional first pw conv.\n    \"\"\"\n    def __init__(self, in_chs, out_chs, dw_kernel_size=3,\n                 stride=1, pad_type='', act_layer=nn.ReLU, noskip=False,\n                 pw_kernel_size=1, pw_act=False, se_ratio=0., se_kwargs=None,\n                 norm_layer=nn.BatchNorm2d, norm_kwargs=None, drop_connect_rate=0.):\n        super(DepthwiseSeparableConv, self).__init__()\n        assert stride in [1, 2]\n        norm_kwargs = norm_kwargs or {}\n        self.has_se = se_ratio is not None and se_ratio > 0.\n        self.has_residual = (stride == 1 and in_chs == out_chs) and not noskip\n        self.drop_connect_rate = drop_connect_rate\n\n        self.conv_dw = select_conv2d(\n            in_chs, in_chs, dw_kernel_size, stride=stride, padding=pad_type, depthwise=True)\n        self.bn1 = norm_layer(in_chs, **norm_kwargs)\n        self.act1 = act_layer(inplace=True)\n\n        # Squeeze-and-excitation\n        if self.has_se:\n            se_kwargs = resolve_se_args(se_kwargs, in_chs, act_layer)\n            self.se = SqueezeExcite(in_chs, se_ratio=se_ratio, **se_kwargs)\n        else:\n            self.se = nn.Identity()\n\n        self.conv_pw = select_conv2d(in_chs, out_chs, pw_kernel_size, padding=pad_type)\n        self.bn2 = norm_layer(out_chs, **norm_kwargs)\n        self.act2 = act_layer(inplace=True) if pw_act else nn.Identity()\n\n    def forward(self, x):\n        residual = x\n\n        x = self.conv_dw(x)\n        x = self.bn1(x)\n        x = self.act1(x)\n\n        x = self.se(x)\n\n        x = self.conv_pw(x)\n        x = self.bn2(x)\n        x = self.act2(x)\n\n        if self.has_residual:\n            if self.drop_connect_rate > 0.:\n                x = drop_connect(x, self.training, self.drop_connect_rate)\n            x += residual\n        return x\n\n\nclass InvertedResidual(nn.Module):\n    \"\"\" Inverted residual block w/ optional SE\"\"\"\n\n    def __init__(self, in_chs, out_chs, dw_kernel_size=3,\n                 stride=1, pad_type='', act_layer=nn.ReLU, noskip=False,\n                 exp_ratio=1.0, exp_kernel_size=1, pw_kernel_size=1,\n                 se_ratio=0., se_kwargs=None, norm_layer=nn.BatchNorm2d, norm_kwargs=None,\n                 conv_kwargs=None, drop_connect_rate=0.):\n        super(InvertedResidual, self).__init__()\n        norm_kwargs = norm_kwargs or {}\n        conv_kwargs = conv_kwargs or {}\n        mid_chs: int = make_divisible(in_chs * exp_ratio)\n        self.has_se = se_ratio is not None and se_ratio > 0.\n        self.has_residual = (in_chs == out_chs and stride == 1) and not noskip\n        self.drop_connect_rate = drop_connect_rate\n\n        # Point-wise expansion\n        self.conv_pw = select_conv2d(in_chs, mid_chs, exp_kernel_size, padding=pad_type, **conv_kwargs)\n        self.bn1 = norm_layer(mid_chs, **norm_kwargs)\n        self.act1 = act_layer(inplace=True)\n\n        # Depth-wise convolution\n        self.conv_dw = select_conv2d(\n            mid_chs, mid_chs, dw_kernel_size, stride=stride, padding=pad_type, depthwise=True, **conv_kwargs)\n        self.bn2 = norm_layer(mid_chs, **norm_kwargs)\n        self.act2 = act_layer(inplace=True)\n\n        # Squeeze-and-excitation\n        if self.has_se:\n            se_kwargs = resolve_se_args(se_kwargs, in_chs, act_layer)\n            self.se = SqueezeExcite(mid_chs, se_ratio=se_ratio, **se_kwargs)\n        else:\n            self.se = nn.Identity()  # for jit.script compat\n\n        # Point-wise linear projection\n        self.conv_pwl = select_conv2d(mid_chs, out_chs, pw_kernel_size, padding=pad_type, **conv_kwargs)\n        self.bn3 = norm_layer(out_chs, **norm_kwargs)\n\n    def forward(self, x):\n        residual = x\n\n        # Point-wise expansion\n        x = self.conv_pw(x)\n        x = self.bn1(x)\n        x = self.act1(x)\n\n        # Depth-wise convolution\n        x = self.conv_dw(x)\n        x = self.bn2(x)\n        x = self.act2(x)\n\n        # Squeeze-and-excitation\n        x = self.se(x)\n\n        # Point-wise linear projection\n        x = self.conv_pwl(x)\n        x = self.bn3(x)\n\n        if self.has_residual:\n            if self.drop_connect_rate > 0.:\n                x = drop_connect(x, self.training, self.drop_connect_rate)\n            x += residual\n        return x\n\n\nclass CondConvResidual(InvertedResidual):\n    \"\"\" Inverted residual block w/ CondConv routing\"\"\"\n\n    def __init__(self, in_chs, out_chs, dw_kernel_size=3,\n                 stride=1, pad_type='', act_layer=nn.ReLU, noskip=False,\n                 exp_ratio=1.0, exp_kernel_size=1, pw_kernel_size=1,\n                 se_ratio=0., se_kwargs=None, norm_layer=nn.BatchNorm2d, norm_kwargs=None,\n                 num_experts=0, drop_connect_rate=0.):\n\n        self.num_experts = num_experts\n        conv_kwargs = dict(num_experts=self.num_experts)\n\n        super(CondConvResidual, self).__init__(\n            in_chs, out_chs, dw_kernel_size=dw_kernel_size, stride=stride, pad_type=pad_type,\n            act_layer=act_layer, noskip=noskip, exp_ratio=exp_ratio, exp_kernel_size=exp_kernel_size,\n            pw_kernel_size=pw_kernel_size, se_ratio=se_ratio, se_kwargs=se_kwargs,\n            norm_layer=norm_layer, norm_kwargs=norm_kwargs, conv_kwargs=conv_kwargs,\n            drop_connect_rate=drop_connect_rate)\n\n        self.routing_fn = nn.Linear(in_chs, self.num_experts)\n\n    def forward(self, x):\n        residual = x\n\n        # CondConv routing\n        pooled_inputs = F.adaptive_avg_pool2d(x, 1).flatten(1)\n        routing_weights = torch.sigmoid(self.routing_fn(pooled_inputs))\n\n        # Point-wise expansion\n        x = self.conv_pw(x, routing_weights)\n        x = self.bn1(x)\n        x = self.act1(x)\n\n        # Depth-wise convolution\n        x = self.conv_dw(x, routing_weights)\n        x = self.bn2(x)\n        x = self.act2(x)\n\n        # Squeeze-and-excitation\n        x = self.se(x)\n\n        # Point-wise linear projection\n        x = self.conv_pwl(x, routing_weights)\n        x = self.bn3(x)\n\n        if self.has_residual:\n            if self.drop_connect_rate > 0.:\n                x = drop_connect(x, self.training, self.drop_connect_rate)\n            x += residual\n        return x\n\n\nclass EdgeResidual(nn.Module):\n    \"\"\" EdgeTPU Residual block with expansion convolution followed by pointwise-linear w/ stride\"\"\"\n\n    def __init__(self, in_chs, out_chs, exp_kernel_size=3, exp_ratio=1.0, fake_in_chs=0,\n                 stride=1, pad_type='', act_layer=nn.ReLU, noskip=False, pw_kernel_size=1,\n                 se_ratio=0., se_kwargs=None, norm_layer=nn.BatchNorm2d, norm_kwargs=None, drop_connect_rate=0.):\n        super(EdgeResidual, self).__init__()\n        norm_kwargs = norm_kwargs or {}\n        mid_chs = make_divisible(fake_in_chs * exp_ratio) if fake_in_chs > 0 else make_divisible(in_chs * exp_ratio)\n        self.has_se = se_ratio is not None and se_ratio > 0.\n        self.has_residual = (in_chs == out_chs and stride == 1) and not noskip\n        self.drop_connect_rate = drop_connect_rate\n\n        # Expansion convolution\n        self.conv_exp = select_conv2d(in_chs, mid_chs, exp_kernel_size, padding=pad_type)\n        self.bn1 = norm_layer(mid_chs, **norm_kwargs)\n        self.act1 = act_layer(inplace=True)\n\n        # Squeeze-and-excitation\n        if self.has_se:\n            se_kwargs = resolve_se_args(se_kwargs, in_chs, act_layer)\n            self.se = SqueezeExcite(mid_chs, se_ratio=se_ratio, **se_kwargs)\n        else:\n            self.se = nn.Identity()\n\n        # Point-wise linear projection\n        self.conv_pwl = select_conv2d(mid_chs, out_chs, pw_kernel_size, stride=stride, padding=pad_type)\n        self.bn2 = nn.BatchNorm2d(out_chs, **norm_kwargs)\n\n    def forward(self, x):\n        residual = x\n\n        # Expansion convolution\n        x = self.conv_exp(x)\n        x = self.bn1(x)\n        x = self.act1(x)\n\n        # Squeeze-and-excitation\n        x = self.se(x)\n\n        # Point-wise linear projection\n        x = self.conv_pwl(x)\n        x = self.bn2(x)\n\n        if self.has_residual:\n            if self.drop_connect_rate > 0.:\n                x = drop_connect(x, self.training, self.drop_connect_rate)\n            x += residual\n\n        return x\n\n\nclass EfficientNetBuilder:\n    \"\"\" Build Trunk Blocks for Efficient/Mobile Networks\n\n    This ended up being somewhat of a cross between\n    https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mnasnet_models.py\n    and\n    https://github.com/facebookresearch/maskrcnn-benchmark/blob/master/maskrcnn_benchmark/modeling/backbone/fbnet_builder.py\n\n    \"\"\"\n\n    def __init__(self, channel_multiplier=1.0, channel_divisor=8, channel_min=None,\n                 pad_type='', act_layer=None, se_kwargs=None,\n                 norm_layer=nn.BatchNorm2d, norm_kwargs=None, drop_connect_rate=0.):\n        self.channel_multiplier = channel_multiplier\n        self.channel_divisor = channel_divisor\n        self.channel_min = channel_min\n        self.pad_type = pad_type\n        self.act_layer = act_layer\n        self.se_kwargs = se_kwargs\n        self.norm_layer = norm_layer\n        self.norm_kwargs = norm_kwargs\n        self.drop_connect_rate = drop_connect_rate\n\n        # updated during build\n        self.in_chs = None\n        self.block_idx = 0\n        self.block_count = 0\n\n    def _round_channels(self, chs):\n        return round_channels(chs, self.channel_multiplier, self.channel_divisor, self.channel_min)\n\n    def _make_block(self, ba):\n        bt = ba.pop('block_type')\n        ba['in_chs'] = self.in_chs\n        ba['out_chs'] = self._round_channels(ba['out_chs'])\n        if 'fake_in_chs' in ba and ba['fake_in_chs']:\n            # FIXME this is a hack to work around mismatch in origin impl input filters for EdgeTPU\n            ba['fake_in_chs'] = self._round_channels(ba['fake_in_chs'])\n        ba['norm_layer'] = self.norm_layer\n        ba['norm_kwargs'] = self.norm_kwargs\n        ba['pad_type'] = self.pad_type\n        # block act fn overrides the model default\n        ba['act_layer'] = ba['act_layer'] if ba['act_layer'] is not None else self.act_layer\n        assert ba['act_layer'] is not None\n        if bt == 'ir':\n            ba['drop_connect_rate'] = self.drop_connect_rate * self.block_idx / self.block_count\n            ba['se_kwargs'] = self.se_kwargs\n            if ba.get('num_experts', 0) > 0:\n                block = CondConvResidual(**ba)\n            else:\n                block = InvertedResidual(**ba)\n        elif bt == 'ds' or bt == 'dsa':\n            ba['drop_connect_rate'] = self.drop_connect_rate * self.block_idx / self.block_count\n            ba['se_kwargs'] = self.se_kwargs\n            block = DepthwiseSeparableConv(**ba)\n        elif bt == 'er':\n            ba['drop_connect_rate'] = self.drop_connect_rate * self.block_idx / self.block_count\n            ba['se_kwargs'] = self.se_kwargs\n            block = EdgeResidual(**ba)\n        elif bt == 'cn':\n            block = ConvBnAct(**ba)\n        else:\n            assert False, 'Uknkown block type (%s) while building model.' % bt\n        self.in_chs = ba['out_chs']  # update in_chs for arg of next block\n        return block\n\n    def _make_stack(self, stack_args):\n        blocks = []\n        # each stack (stage) contains a list of block arguments\n        for i, ba in enumerate(stack_args):\n            if i >= 1:\n                # only the first block in any stack can have a stride > 1\n                ba['stride'] = 1\n            block = self._make_block(ba)\n            blocks.append(block)\n            self.block_idx += 1  # incr global idx (across all stacks)\n        return nn.Sequential(*blocks)\n\n    def __call__(self, in_chs, block_args):\n        \"\"\" Build the blocks\n        Args:\n            in_chs: Number of input-channels passed to first block\n            block_args: A list of lists, outer list defines stages, inner\n                list contains strings defining block configuration(s)\n        Return:\n             List of block stacks (each stack wrapped in nn.Sequential)\n        \"\"\"\n        self.in_chs = in_chs\n        self.block_count = sum([len(x) for x in block_args])\n        self.block_idx = 0\n        blocks = []\n        # outer list of block_args defines the stacks ('stages' by some conventions)\n        for stack_idx, stack in enumerate(block_args):\n            assert isinstance(stack, list)\n            stack = self._make_stack(stack)\n            blocks.append(stack)\n        return blocks\n\n\ndef _parse_ksize(ss):\n    if ss.isdigit():\n        return int(ss)\n    else:\n        return [int(k) for k in ss.split('.')]\n\n\ndef _decode_block_str(block_str):\n    \"\"\" Decode block definition string\n\n    Gets a list of block arg (dicts) through a string notation of arguments.\n    E.g. ir_r2_k3_s2_e1_i32_o16_se0.25_noskip\n\n    All args can exist in any order with the exception of the leading string which\n    is assumed to indicate the block type.\n\n    leading string - block type (\n      ir = InvertedResidual, ds = DepthwiseSep, dsa = DeptwhiseSep with pw act, cn = ConvBnAct)\n    r - number of repeat blocks,\n    k - kernel size,\n    s - strides (1-9),\n    e - expansion ratio,\n    c - output channels,\n    se - squeeze/excitation ratio\n    n - activation fn ('re', 'r6', 'hs', or 'sw')\n    Args:\n        block_str: a string representation of block arguments.\n    Returns:\n        A list of block args (dicts)\n    Raises:\n        ValueError: if the string def not properly specified (TODO)\n    \"\"\"\n    assert isinstance(block_str, str)\n    ops = block_str.split('_')\n    block_type = ops[0]  # take the block type off the front\n    ops = ops[1:]\n    options = {}\n    noskip = False\n    for op in ops:\n        # string options being checked on individual basis, combine if they grow\n        if op == 'noskip':\n            noskip = True\n        elif op.startswith('n'):\n            # activation fn\n            key = op[0]\n            v = op[1:]\n            if v == 're':\n                value = get_act_layer('relu')\n            elif v == 'r6':\n                value = get_act_layer('relu6')\n            elif v == 'hs':\n                value = get_act_layer('hard_swish')\n            elif v == 'sw':\n                value = get_act_layer('swish')\n            else:\n                continue\n            options[key] = value\n        else:\n            # all numeric options\n            splits = re.split(r'(\\d.*)', op)\n            if len(splits) >= 2:\n                key, value = splits[:2]\n                options[key] = value\n\n    # if act_layer is None, the model default (passed to model init) will be used\n    act_layer = options['n'] if 'n' in options else None\n    exp_kernel_size = _parse_ksize(options['a']) if 'a' in options else 1\n    pw_kernel_size = _parse_ksize(options['p']) if 'p' in options else 1\n    fake_in_chs = int(options['fc']) if 'fc' in options else 0  # FIXME hack to deal with in_chs issue in TPU def\n\n    num_repeat = int(options['r'])\n    # each type of block has different valid arguments, fill accordingly\n    if block_type == 'ir':\n        block_args = dict(\n            block_type=block_type,\n            dw_kernel_size=_parse_ksize(options['k']),\n            exp_kernel_size=exp_kernel_size,\n            pw_kernel_size=pw_kernel_size,\n            out_chs=int(options['c']),\n            exp_ratio=float(options['e']),\n            se_ratio=float(options['se']) if 'se' in options else None,\n            stride=int(options['s']),\n            act_layer=act_layer,\n            noskip=noskip,\n        )\n        if 'cc' in options:\n            block_args['num_experts'] = int(options['cc'])\n    elif block_type == 'ds' or block_type == 'dsa':\n        block_args = dict(\n            block_type=block_type,\n            dw_kernel_size=_parse_ksize(options['k']),\n            pw_kernel_size=pw_kernel_size,\n            out_chs=int(options['c']),\n            se_ratio=float(options['se']) if 'se' in options else None,\n            stride=int(options['s']),\n            act_layer=act_layer,\n            pw_act=block_type == 'dsa',\n            noskip=block_type == 'dsa' or noskip,\n        )\n    elif block_type == 'er':\n        block_args = dict(\n            block_type=block_type,\n            exp_kernel_size=_parse_ksize(options['k']),\n            pw_kernel_size=pw_kernel_size,\n            out_chs=int(options['c']),\n            exp_ratio=float(options['e']),\n            fake_in_chs=fake_in_chs,\n            se_ratio=float(options['se']) if 'se' in options else None,\n            stride=int(options['s']),\n            act_layer=act_layer,\n            noskip=noskip,\n        )\n    elif block_type == 'cn':\n        block_args = dict(\n            block_type=block_type,\n            kernel_size=int(options['k']),\n            out_chs=int(options['c']),\n            stride=int(options['s']),\n            act_layer=act_layer,\n        )\n    else:\n        assert False, 'Unknown block type (%s)' % block_type\n\n    return block_args, num_repeat\n\n\ndef _scale_stage_depth(stack_args, repeats, depth_multiplier=1.0, depth_trunc='ceil'):\n    \"\"\" Per-stage depth scaling\n    Scales the block repeats in each stage. This depth scaling impl maintains\n    compatibility with the EfficientNet scaling method, while allowing sensible\n    scaling for other models that may have multiple block arg definitions in each stage.\n    \"\"\"\n\n    # We scale the total repeat count for each stage, there may be multiple\n    # block arg defs per stage so we need to sum.\n    num_repeat = sum(repeats)\n    if depth_trunc == 'round':\n        # Truncating to int by rounding allows stages with few repeats to remain\n        # proportionally smaller for longer. This is a good choice when stage definitions\n        # include single repeat stages that we'd prefer to keep that way as long as possible\n        num_repeat_scaled = max(1, round(num_repeat * depth_multiplier))\n    else:\n        # The default for EfficientNet truncates repeats to int via 'ceil'.\n        # Any multiplier > 1.0 will result in an increased depth for every stage.\n        num_repeat_scaled = int(math.ceil(num_repeat * depth_multiplier))\n\n    # Proportionally distribute repeat count scaling to each block definition in the stage.\n    # Allocation is done in reverse as it results in the first block being less likely to be scaled.\n    # The first block makes less sense to repeat in most of the arch definitions.\n    repeats_scaled = []\n    for r in repeats[::-1]:\n        rs = max(1, round((r / num_repeat * num_repeat_scaled)))\n        repeats_scaled.append(rs)\n        num_repeat -= r\n        num_repeat_scaled -= rs\n    repeats_scaled = repeats_scaled[::-1]\n\n    # Apply the calculated scaling to each block arg in the stage\n    sa_scaled = []\n    for ba, rep in zip(stack_args, repeats_scaled):\n        sa_scaled.extend([deepcopy(ba) for _ in range(rep)])\n    return sa_scaled\n\n\ndef decode_arch_def(arch_def, depth_multiplier=1.0, depth_trunc='ceil', experts_multiplier=1):\n    arch_args = []\n    for stack_idx, block_strings in enumerate(arch_def):\n        assert isinstance(block_strings, list)\n        stack_args = []\n        repeats = []\n        for block_str in block_strings:\n            assert isinstance(block_str, str)\n            ba, rep = _decode_block_str(block_str)\n            if ba.get('num_experts', 0) > 0 and experts_multiplier > 1:\n                ba['num_experts'] *= experts_multiplier\n            stack_args.append(ba)\n            repeats.append(rep)\n        arch_args.append(_scale_stage_depth(stack_args, repeats, depth_multiplier, depth_trunc))\n    return arch_args\n\n\ndef initialize_weight_goog(m, n=''):\n    # weight init as per Tensorflow Official impl\n    # https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mnasnet_model.py\n    if isinstance(m, CondConv2d):\n        fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n        init_weight_fn = get_condconv_initializer(\n            lambda w: w.data.normal_(0, math.sqrt(2.0 / fan_out)), m.num_experts, m.weight_shape)\n        init_weight_fn(m.weight)\n        if m.bias is not None:\n            m.bias.data.zero_()\n    elif isinstance(m, nn.Conv2d):\n        fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n        m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))\n        if m.bias is not None:\n            m.bias.data.zero_()\n    elif isinstance(m, nn.BatchNorm2d):\n        m.weight.data.fill_(1.0)\n        m.bias.data.zero_()\n    elif isinstance(m, nn.Linear):\n        fan_out = m.weight.size(0)  # fan-out\n        fan_in = 0\n        if 'routing_fn' in n:\n            fan_in = m.weight.size(1)\n        init_range = 1.0 / math.sqrt(fan_in + fan_out)\n        m.weight.data.uniform_(-init_range, init_range)\n        m.bias.data.zero_()\n\n\ndef initialize_weight_default(m, n=''):\n    if isinstance(m, CondConv2d):\n        init_fn = get_condconv_initializer(partial(\n            nn.init.kaiming_normal_, mode='fan_out', nonlinearity='relu'), m.num_experts, m.weight_shape)\n        init_fn(m.weight)\n    elif isinstance(m, nn.Conv2d):\n        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n    elif isinstance(m, nn.BatchNorm2d):\n        m.weight.data.fill_(1.0)\n        m.bias.data.zero_()\n    elif isinstance(m, nn.Linear):\n        nn.init.kaiming_uniform_(m.weight, mode='fan_in', nonlinearity='linear')\n"
  },
  {
    "path": "backbones/geffnet/gen_efficientnet.py",
    "content": "\"\"\" Generic Efficient Networks\n\nA generic MobileNet class with building blocks to support a variety of models:\n\n* EfficientNet (B0-B8 + Tensorflow pretrained AutoAug/RandAug/AdvProp ports)\n  - EfficientNet: Rethinking Model Scaling for CNNs - https://arxiv.org/abs/1905.11946\n  - CondConv: Conditionally Parameterized Convolutions for Efficient Inference - https://arxiv.org/abs/1904.04971\n  - Adversarial Examples Improve Image Recognition - https://arxiv.org/abs/1911.09665\n\n* MixNet (Small, Medium, and Large)\n  - MixConv: Mixed Depthwise Convolutional Kernels - https://arxiv.org/abs/1907.09595\n\n* MNasNet B1, A1 (SE), Small\n  - MnasNet: Platform-Aware Neural Architecture Search for Mobile - https://arxiv.org/abs/1807.11626\n\n* FBNet-C\n  - FBNet: Hardware-Aware Efficient ConvNet Design via Differentiable NAS - https://arxiv.org/abs/1812.03443\n\n* Single-Path NAS Pixel1\n  - Single-Path NAS: Designing Hardware-Efficient ConvNets - https://arxiv.org/abs/1904.02877\n\n* And likely more...\n\nHacked together by Ross Wightman\n\"\"\"\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom .helpers import load_pretrained\nfrom .efficientnet_builder import *\n\n__all__ = ['GenEfficientNet', 'mnasnet_050', 'mnasnet_075', 'mnasnet_100', 'mnasnet_b1', 'mnasnet_140',\n           'semnasnet_050', 'semnasnet_075', 'semnasnet_100', 'mnasnet_a1', 'semnasnet_140', 'mnasnet_small',\n           'fbnetc_100', 'spnasnet_100', 'efficientnet_b0', 'efficientnet_b1', 'efficientnet_b2',  'efficientnet_b3',\n           'efficientnet_b4', 'efficientnet_b5', 'efficientnet_b6', 'efficientnet_b7', 'efficientnet_b8',\n           'efficientnet_es', 'efficientnet_em', 'efficientnet_el',\n           'efficientnet_cc_b0_4e', 'efficientnet_cc_b0_8e', 'efficientnet_cc_b1_8e',\n           'tf_efficientnet_b0', 'tf_efficientnet_b1', 'tf_efficientnet_b2', 'tf_efficientnet_b3',\n           'tf_efficientnet_b4', 'tf_efficientnet_b5', 'tf_efficientnet_b6', 'tf_efficientnet_b7',\n           'tf_efficientnet_b0_ap', 'tf_efficientnet_b1_ap', 'tf_efficientnet_b2_ap', 'tf_efficientnet_b3_ap',\n           'tf_efficientnet_b4_ap', 'tf_efficientnet_b5_ap', 'tf_efficientnet_b6_ap', 'tf_efficientnet_b7_ap',\n           'tf_efficientnet_b8_ap', 'tf_efficientnet_es', 'tf_efficientnet_em', 'tf_efficientnet_el',\n           'tf_efficientnet_cc_b0_4e', 'tf_efficientnet_cc_b0_8e', 'tf_efficientnet_cc_b1_8e',\n           'mixnet_s', 'mixnet_m', 'mixnet_l', 'mixnet_xl', 'tf_mixnet_s', 'tf_mixnet_m', 'tf_mixnet_l']\n\n\nmodel_urls = {\n    'mnasnet_050': None,\n    'mnasnet_075': None,\n    'mnasnet_100':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mnasnet_b1-74cb7081.pth',\n    'mnasnet_140': None,\n    'semnasnet_050': None,\n    'semnasnet_075': None,\n    'semnasnet_100':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mnasnet_a1-d9418771.pth',\n    'semnasnet_140': None,\n    'mnasnet_small': None,\n    'fbnetc_100':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/fbnetc_100-c345b898.pth',\n    'spnasnet_100':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/spnasnet_100-048bc3f4.pth',\n    'efficientnet_b0':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/efficientnet_b0-d6904d92.pth',\n    'efficientnet_b1':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/efficientnet_b1-533bc792.pth',\n    'efficientnet_b2':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/efficientnet_b2-cf78dc4d.pth',\n    'efficientnet_b3': None,\n    'efficientnet_b4': None,\n    'efficientnet_b5': None,\n    'efficientnet_b6': None,\n    'efficientnet_b7': None,\n    'efficientnet_b8': None,\n    'efficientnet_es': None,\n    'efficientnet_em': None,\n    'efficientnet_el': None,\n    'efficientnet_cc_b0_4e': None,\n    'efficientnet_cc_b0_8e': None,\n    'efficientnet_cc_b1_8e': None,\n    'tf_efficientnet_b0':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b0_aa-827b6e33.pth',\n    'tf_efficientnet_b1':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b1_aa-ea7a6ee0.pth',\n    'tf_efficientnet_b2':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b2_aa-60c94f97.pth',\n    'tf_efficientnet_b3':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b3_aa-84b4657e.pth',\n    'tf_efficientnet_b4':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b4_aa-818f208c.pth',\n    'tf_efficientnet_b5':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b5_ra-9a3e5369.pth',\n    'tf_efficientnet_b6':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b6_aa-80ba17e4.pth',\n    'tf_efficientnet_b7':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b7_ra-6c08e654.pth',\n    'tf_efficientnet_b0_ap':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b0_ap-f262efe1.pth',\n    'tf_efficientnet_b1_ap':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b1_ap-44ef0a3d.pth',\n    'tf_efficientnet_b2_ap':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b2_ap-2f8e7636.pth',\n    'tf_efficientnet_b3_ap':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b3_ap-aad25bdd.pth',\n    'tf_efficientnet_b4_ap':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b4_ap-dedb23e6.pth',\n    'tf_efficientnet_b5_ap':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b5_ap-9e82fae8.pth',\n    'tf_efficientnet_b6_ap':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b6_ap-4ffb161f.pth',\n    'tf_efficientnet_b7_ap':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b7_ap-ddb28fec.pth',\n    'tf_efficientnet_b8_ap':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b8_ap-00e169fa.pth',\n    'tf_efficientnet_es':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_es-ca1afbfe.pth',\n    'tf_efficientnet_em':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_em-e78cfe58.pth',\n    'tf_efficientnet_el':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_el-5143854e.pth',\n    'tf_efficientnet_cc_b0_4e':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_cc_b0_4e-4362b6b2.pth',\n    'tf_efficientnet_cc_b0_8e':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_cc_b0_8e-66184a25.pth',\n    'tf_efficientnet_cc_b1_8e':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_cc_b1_8e-f7c79ae1.pth',\n    'mixnet_s': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mixnet_s-a907afbc.pth',\n    'mixnet_m': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mixnet_m-4647fc68.pth',\n    'mixnet_l': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mixnet_l-5a9a2ed8.pth',\n    'mixnet_xl': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mixnet_xl-ac5fbe8d.pth',\n    'tf_mixnet_s':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mixnet_s-89d3354b.pth',\n    'tf_mixnet_m':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mixnet_m-0f4d8805.pth',\n    'tf_mixnet_l':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mixnet_l-6c92e0c8.pth',\n}\n\n\nclass GenEfficientNet(nn.Module):\n    \"\"\" Generic EfficientNets\n\n    An implementation of mobile optimized networks that covers:\n      * EfficientNet (B0-B8, CondConv, EdgeTPU)\n      * MixNet (Small, Medium, and Large, XL)\n      * MNASNet A1, B1, and small\n      * FBNet C\n      * Single-Path NAS Pixel1\n    \"\"\"\n\n    def __init__(self, block_args, num_classes=1000, in_chans=3, stem_size=32, num_features=1280,\n                 channel_multiplier=1.0, channel_divisor=8, channel_min=None,\n                 pad_type='', act_layer=nn.ReLU, drop_rate=0., drop_connect_rate=0.,\n                 se_kwargs=None, norm_layer=nn.BatchNorm2d, norm_kwargs=None,\n                 weight_init='goog'):\n        super(GenEfficientNet, self).__init__()\n        self.drop_rate = drop_rate\n\n        stem_size = round_channels(stem_size, channel_multiplier, channel_divisor, channel_min)\n        self.conv_stem = select_conv2d(in_chans, stem_size, 3, stride=2, padding=pad_type)\n        self.bn1 = norm_layer(stem_size, **norm_kwargs)\n        self.act1 = act_layer(inplace=True)\n        in_chs = stem_size\n\n        builder = EfficientNetBuilder(\n            channel_multiplier, channel_divisor, channel_min,\n            pad_type, act_layer, se_kwargs, norm_layer, norm_kwargs, drop_connect_rate)\n        # modulelist = builder(in_chs, block_args)\n        self.blocks = nn.ModuleList()\n        for block in builder(in_chs,block_args):\n            self.blocks.append(block)\n\n        in_chs = builder.in_chs\n\n        self.conv_head = select_conv2d(in_chs, num_features, 1, padding=pad_type)\n        self.bn2 = norm_layer(num_features, **norm_kwargs)\n        self.act2 = act_layer(inplace=True)\n        self.global_pool = nn.AdaptiveAvgPool2d(1)\n        self.classifier = nn.Linear(num_features, num_classes)\n\n        for n, m in self.named_modules():\n            if weight_init == 'goog':\n                initialize_weight_goog(m, n)\n            else:\n                initialize_weight_default(m, n)\n\n    def features(self, x):\n        x = self.conv_stem(x)\n        x = self.bn1(x)\n        x = self.act1(x)\n        outs=[]\n        for block in self.blocks:\n            x = block(x)\n            outs.append(x)\n        # x = self.blocks(x)\n        # x = self.conv_head(x)\n        # x = self.bn2(x)\n        # x = self.act2(x)\n        return outs\n\n    def as_sequential(self):\n        layers = [self.conv_stem, self.bn1, self.act1]\n        layers.extend(self.blocks)\n        layers.extend([\n            self.conv_head, self.bn2, self.act2,\n            self.global_pool, nn.Flatten(), nn.Dropout(self.drop_rate), self.classifier])\n        return nn.Sequential(*layers)\n\n    def forward(self, x):\n        # x = self.features(x)\n        # x = self.global_pool(x)\n        # x = x.flatten(1)\n        # if self.drop_rate > 0.:\n        #     x = F.dropout(x, p=self.drop_rate, training=self.training)\n        return self.features(x)\n\ndef _create_model(model_kwargs, variant, pretrained=False):\n    as_sequential = model_kwargs.pop('as_sequential', False)\n    model = GenEfficientNet(**model_kwargs)\n    if pretrained and model_urls[variant]:\n        load_pretrained(model, model_urls[variant])\n    if as_sequential:\n        model = model.as_sequential()\n    return model\n\n\ndef _gen_mnasnet_a1(variant, channel_multiplier=1.0, pretrained=False, **kwargs):\n    \"\"\"Creates a mnasnet-a1 model.\n\n    Ref impl: https://github.com/tensorflow/tpu/tree/master/models/official/mnasnet\n    Paper: https://arxiv.org/pdf/1807.11626.pdf.\n\n    Args:\n      channel_multiplier: multiplier to number of channels per layer.\n    \"\"\"\n    arch_def = [\n        # stage 0, 112x112 in\n        ['ds_r1_k3_s1_e1_c16_noskip'],\n        # stage 1, 112x112 in\n        ['ir_r2_k3_s2_e6_c24'],\n        # stage 2, 56x56 in\n        ['ir_r3_k5_s2_e3_c40_se0.25'],\n        # stage 3, 28x28 in\n        ['ir_r4_k3_s2_e6_c80'],\n        # stage 4, 14x14in\n        ['ir_r2_k3_s1_e6_c112_se0.25'],\n        # stage 5, 14x14in\n        ['ir_r3_k5_s2_e6_c160_se0.25'],\n        # stage 6, 7x7 in\n        ['ir_r1_k3_s1_e6_c320'],\n    ]\n    model_kwargs = dict(\n        block_args=decode_arch_def(arch_def),\n        stem_size=32,\n        channel_multiplier=channel_multiplier,\n        act_layer=resolve_act_layer(kwargs, 'relu'),\n        norm_kwargs=resolve_bn_args(kwargs),\n        **kwargs\n    )\n    model = _create_model(model_kwargs, variant, pretrained)\n    return model\n\n\ndef _gen_mnasnet_b1(variant, channel_multiplier=1.0, pretrained=False, **kwargs):\n    \"\"\"Creates a mnasnet-b1 model.\n\n    Ref impl: https://github.com/tensorflow/tpu/tree/master/models/official/mnasnet\n    Paper: https://arxiv.org/pdf/1807.11626.pdf.\n\n    Args:\n      channel_multiplier: multiplier to number of channels per layer.\n    \"\"\"\n    arch_def = [\n        # stage 0, 112x112 in\n        ['ds_r1_k3_s1_c16_noskip'],\n        # stage 1, 112x112 in\n        ['ir_r3_k3_s2_e3_c24'],\n        # stage 2, 56x56 in\n        ['ir_r3_k5_s2_e3_c40'],\n        # stage 3, 28x28 in\n        ['ir_r3_k5_s2_e6_c80'],\n        # stage 4, 14x14in\n        ['ir_r2_k3_s1_e6_c96'],\n        # stage 5, 14x14in\n        ['ir_r4_k5_s2_e6_c192'],\n        # stage 6, 7x7 in\n        ['ir_r1_k3_s1_e6_c320_noskip']\n    ]\n    model_kwargs = dict(\n        block_args=decode_arch_def(arch_def),\n        stem_size=32,\n        channel_multiplier=channel_multiplier,\n        act_layer=resolve_act_layer(kwargs, 'relu'),\n        norm_kwargs=resolve_bn_args(kwargs),\n        **kwargs\n    )\n    model = _create_model(model_kwargs, variant, pretrained)\n    return model\n\n\ndef _gen_mnasnet_small(variant, channel_multiplier=1.0, pretrained=False, **kwargs):\n    \"\"\"Creates a mnasnet-b1 model.\n\n    Ref impl: https://github.com/tensorflow/tpu/tree/master/models/official/mnasnet\n    Paper: https://arxiv.org/pdf/1807.11626.pdf.\n\n    Args:\n      channel_multiplier: multiplier to number of channels per layer.\n    \"\"\"\n    arch_def = [\n        ['ds_r1_k3_s1_c8'],\n        ['ir_r1_k3_s2_e3_c16'],\n        ['ir_r2_k3_s2_e6_c16'],\n        ['ir_r4_k5_s2_e6_c32_se0.25'],\n        ['ir_r3_k3_s1_e6_c32_se0.25'],\n        ['ir_r3_k5_s2_e6_c88_se0.25'],\n        ['ir_r1_k3_s1_e6_c144']\n    ]\n    model_kwargs = dict(\n        block_args=decode_arch_def(arch_def),\n        stem_size=8,\n        channel_multiplier=channel_multiplier,\n        act_layer=resolve_act_layer(kwargs, 'relu'),\n        norm_kwargs=resolve_bn_args(kwargs),\n        **kwargs\n    )\n    model = _create_model(model_kwargs, variant, pretrained)\n    return model\n\n\ndef _gen_fbnetc(variant, channel_multiplier=1.0, pretrained=False, **kwargs):\n    \"\"\" FBNet-C\n\n        Paper: https://arxiv.org/abs/1812.03443\n        Ref Impl: https://github.com/facebookresearch/maskrcnn-benchmark/blob/master/maskrcnn_benchmark/modeling/backbone/fbnet_modeldef.py\n\n        NOTE: the impl above does not relate to the 'C' variant here, that was derived from paper,\n        it was used to confirm some building block details\n    \"\"\"\n    arch_def = [\n        ['ir_r1_k3_s1_e1_c16'],\n        ['ir_r1_k3_s2_e6_c24', 'ir_r2_k3_s1_e1_c24'],\n        ['ir_r1_k5_s2_e6_c32', 'ir_r1_k5_s1_e3_c32', 'ir_r1_k5_s1_e6_c32', 'ir_r1_k3_s1_e6_c32'],\n        ['ir_r1_k5_s2_e6_c64', 'ir_r1_k5_s1_e3_c64', 'ir_r2_k5_s1_e6_c64'],\n        ['ir_r3_k5_s1_e6_c112', 'ir_r1_k5_s1_e3_c112'],\n        ['ir_r4_k5_s2_e6_c184'],\n        ['ir_r1_k3_s1_e6_c352'],\n    ]\n    model_kwargs = dict(\n        block_args=decode_arch_def(arch_def),\n        stem_size=16,\n        num_features=1984,  # paper suggests this, but is not 100% clear\n        channel_multiplier=channel_multiplier,\n        act_layer=resolve_act_layer(kwargs, 'relu'),\n        norm_kwargs=resolve_bn_args(kwargs),\n        **kwargs\n    )\n    model = _create_model(model_kwargs, variant, pretrained)\n    return model\n\n\ndef _gen_spnasnet(variant, channel_multiplier=1.0, pretrained=False, **kwargs):\n    \"\"\"Creates the Single-Path NAS model from search targeted for Pixel1 phone.\n\n    Paper: https://arxiv.org/abs/1904.02877\n\n    Args:\n      channel_multiplier: multiplier to number of channels per layer.\n    \"\"\"\n    arch_def = [\n        # stage 0, 112x112 in\n        ['ds_r1_k3_s1_c16_noskip'],\n        # stage 1, 112x112 in\n        ['ir_r3_k3_s2_e3_c24'],\n        # stage 2, 56x56 in\n        ['ir_r1_k5_s2_e6_c40', 'ir_r3_k3_s1_e3_c40'],\n        # stage 3, 28x28 in\n        ['ir_r1_k5_s2_e6_c80', 'ir_r3_k3_s1_e3_c80'],\n        # stage 4, 14x14in\n        ['ir_r1_k5_s1_e6_c96', 'ir_r3_k5_s1_e3_c96'],\n        # stage 5, 14x14in\n        ['ir_r4_k5_s2_e6_c192'],\n        # stage 6, 7x7 in\n        ['ir_r1_k3_s1_e6_c320_noskip']\n    ]\n    model_kwargs = dict(\n        block_args=decode_arch_def(arch_def),\n        stem_size=32,\n        channel_multiplier=channel_multiplier,\n        act_layer=resolve_act_layer(kwargs, 'relu'),\n        norm_kwargs=resolve_bn_args(kwargs),\n        **kwargs\n    )\n    model = _create_model(model_kwargs, variant, pretrained)\n    return model\n\n\ndef _gen_efficientnet(variant, channel_multiplier=1.0, depth_multiplier=1.0, pretrained=False, **kwargs):\n    \"\"\"Creates an EfficientNet model.\n\n    Ref impl: https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/efficientnet_model.py\n    Paper: https://arxiv.org/abs/1905.11946\n\n    EfficientNet params\n    name: (channel_multiplier, depth_multiplier, resolution, dropout_rate)\n    'efficientnet-b0': (1.0, 1.0, 224, 0.2),\n    'efficientnet-b1': (1.0, 1.1, 240, 0.2),\n    'efficientnet-b2': (1.1, 1.2, 260, 0.3),\n    'efficientnet-b3': (1.2, 1.4, 300, 0.3),\n    'efficientnet-b4': (1.4, 1.8, 380, 0.4),\n    'efficientnet-b5': (1.6, 2.2, 456, 0.4),\n    'efficientnet-b6': (1.8, 2.6, 528, 0.5),\n    'efficientnet-b7': (2.0, 3.1, 600, 0.5),\n    'efficientnet-b8': (2.2, 3.6, 672, 0.5),\n\n    Args:\n      channel_multiplier: multiplier to number of channels per layer\n      depth_multiplier: multiplier to number of repeats per stage\n\n    \"\"\"\n    # s1->s2\n    arch_def = [\n        ['ds_r1_k3_s1_e1_c16_se0.25'],\n        ['ir_r2_k3_s2_e6_c24_se0.25'],\n        ['ir_r2_k5_s2_e6_c40_se0.25'],\n        ['ir_r3_k3_s2_e6_c80_se0.25'],\n        ['ir_r3_k5_s2_e6_c112_se0.25'],\n        ['ir_r4_k5_s2_e6_c192_se0.25'],\n        ['ir_r1_k3_s2_e6_c320_se0.25'],\n    ]\n    model_kwargs = dict(\n        block_args=decode_arch_def(arch_def, depth_multiplier),\n        num_features=round_channels(1280, channel_multiplier, 8, None),\n        stem_size=32,\n        channel_multiplier=channel_multiplier,\n        act_layer=resolve_act_layer(kwargs, 'swish'),\n        norm_kwargs=resolve_bn_args(kwargs),\n        **kwargs,\n    )\n    model = _create_model(model_kwargs, variant, pretrained)\n    return model\n\n\ndef _gen_efficientnet_edge(variant, channel_multiplier=1.0, depth_multiplier=1.0, pretrained=False, **kwargs):\n    arch_def = [\n        # NOTE `fc` is present to override a mismatch between stem channels and in chs not\n        # present in other models\n        ['er_r1_k3_s1_e4_c24_fc24_noskip'],\n        ['er_r2_k3_s2_e8_c32'],\n        ['er_r4_k3_s2_e8_c48'],\n        ['ir_r5_k5_s2_e8_c96'],\n        ['ir_r4_k5_s1_e8_c144'],\n        ['ir_r2_k5_s2_e8_c192'],\n    ]\n    model_kwargs = dict(\n        block_args=decode_arch_def(arch_def, depth_multiplier),\n        num_features=round_channels(1280, channel_multiplier, 8, None),\n        stem_size=32,\n        channel_multiplier=channel_multiplier,\n        act_layer=resolve_act_layer(kwargs, 'relu'),\n        norm_kwargs=resolve_bn_args(kwargs),\n        **kwargs,\n    )\n    model = _create_model(model_kwargs, variant, pretrained)\n    return model\n\n\ndef _gen_efficientnet_condconv(\n        variant, channel_multiplier=1.0, depth_multiplier=1.0, experts_multiplier=1, pretrained=False, **kwargs):\n    \"\"\"Creates an efficientnet-condconv model.\"\"\"\n    arch_def = [\n      ['ds_r1_k3_s1_e1_c16_se0.25'],\n      ['ir_r2_k3_s2_e6_c24_se0.25'],\n      ['ir_r2_k5_s2_e6_c40_se0.25'],\n      ['ir_r3_k3_s2_e6_c80_se0.25'],\n      ['ir_r3_k5_s1_e6_c112_se0.25_cc4'],\n      ['ir_r4_k5_s2_e6_c192_se0.25_cc4'],\n      ['ir_r1_k3_s1_e6_c320_se0.25_cc4'],\n    ]\n    model_kwargs = dict(\n        block_args=decode_arch_def(arch_def, depth_multiplier, experts_multiplier=experts_multiplier),\n        num_features=round_channels(1280, channel_multiplier, 8, None),\n        stem_size=32,\n        channel_multiplier=channel_multiplier,\n        act_layer=resolve_act_layer(kwargs, 'swish'),\n        norm_kwargs=resolve_bn_args(kwargs),\n        **kwargs,\n    )\n    model = _create_model(model_kwargs, variant, pretrained)\n    return model\n\n\ndef _gen_mixnet_s(variant, channel_multiplier=1.0, pretrained=False, **kwargs):\n    \"\"\"Creates a MixNet Small model.\n\n    Ref impl: https://github.com/tensorflow/tpu/tree/master/models/official/mnasnet/mixnet\n    Paper: https://arxiv.org/abs/1907.09595\n    \"\"\"\n    arch_def = [\n        # stage 0, 112x112 in\n        ['ds_r1_k3_s1_e1_c16'],  # relu\n        # stage 1, 112x112 in\n        ['ir_r1_k3_a1.1_p1.1_s2_e6_c24', 'ir_r1_k3_a1.1_p1.1_s1_e3_c24'],  # relu\n        # stage 2, 56x56 in\n        ['ir_r1_k3.5.7_s2_e6_c40_se0.5_nsw', 'ir_r3_k3.5_a1.1_p1.1_s1_e6_c40_se0.5_nsw'],  # swish\n        # stage 3, 28x28 in\n        ['ir_r1_k3.5.7_p1.1_s2_e6_c80_se0.25_nsw', 'ir_r2_k3.5_p1.1_s1_e6_c80_se0.25_nsw'],  # swish\n        # stage 4, 14x14in\n        ['ir_r1_k3.5.7_a1.1_p1.1_s1_e6_c120_se0.5_nsw', 'ir_r2_k3.5.7.9_a1.1_p1.1_s1_e3_c120_se0.5_nsw'],  # swish\n        # stage 5, 14x14in\n        ['ir_r1_k3.5.7.9.11_s2_e6_c200_se0.5_nsw', 'ir_r2_k3.5.7.9_p1.1_s1_e6_c200_se0.5_nsw'],  # swish\n        # 7x7\n    ]\n    model_kwargs = dict(\n        block_args=decode_arch_def(arch_def),\n        num_features=1536,\n        stem_size=16,\n        channel_multiplier=channel_multiplier,\n        act_layer=resolve_act_layer(kwargs, 'relu'),\n        norm_kwargs=resolve_bn_args(kwargs),\n        **kwargs\n    )\n    model = _create_model(model_kwargs, variant, pretrained)\n    return model\n\n\ndef _gen_mixnet_m(variant, channel_multiplier=1.0, depth_multiplier=1.0, pretrained=False, **kwargs):\n    \"\"\"Creates a MixNet Medium-Large model.\n\n    Ref impl: https://github.com/tensorflow/tpu/tree/master/models/official/mnasnet/mixnet\n    Paper: https://arxiv.org/abs/1907.09595\n    \"\"\"\n    arch_def = [\n        # stage 0, 112x112 in\n        ['ds_r1_k3_s1_e1_c24'],  # relu\n        # stage 1, 112x112 in\n        ['ir_r1_k3.5.7_a1.1_p1.1_s2_e6_c32', 'ir_r1_k3_a1.1_p1.1_s1_e3_c32'],  # relu\n        # stage 2, 56x56 in\n        ['ir_r1_k3.5.7.9_s2_e6_c40_se0.5_nsw', 'ir_r3_k3.5_a1.1_p1.1_s1_e6_c40_se0.5_nsw'],  # swish\n        # stage 3, 28x28 in\n        ['ir_r1_k3.5.7_s2_e6_c80_se0.25_nsw', 'ir_r3_k3.5.7.9_a1.1_p1.1_s1_e6_c80_se0.25_nsw'],  # swish\n        # stage 4, 14x14in\n        ['ir_r1_k3_s1_e6_c120_se0.5_nsw', 'ir_r3_k3.5.7.9_a1.1_p1.1_s1_e3_c120_se0.5_nsw'],  # swish\n        # stage 5, 14x14in\n        ['ir_r1_k3.5.7.9_s2_e6_c200_se0.5_nsw', 'ir_r3_k3.5.7.9_p1.1_s1_e6_c200_se0.5_nsw'],  # swish\n        # 7x7\n    ]\n    model_kwargs = dict(\n        block_args=decode_arch_def(arch_def, depth_multiplier, depth_trunc='round'),\n        num_features=1536,\n        stem_size=24,\n        channel_multiplier=channel_multiplier,\n        act_layer=resolve_act_layer(kwargs, 'relu'),\n        norm_kwargs=resolve_bn_args(kwargs),\n        **kwargs\n    )\n    model = _create_model(model_kwargs, variant, pretrained)\n    return model\n\n\ndef mnasnet_050(pretrained=False, **kwargs):\n    \"\"\" MNASNet B1, depth multiplier of 0.5. \"\"\"\n    model = _gen_mnasnet_b1('mnasnet_050', 0.5, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mnasnet_075(pretrained=False, **kwargs):\n    \"\"\" MNASNet B1, depth multiplier of 0.75. \"\"\"\n    model = _gen_mnasnet_b1('mnasnet_075', 0.75, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mnasnet_100(pretrained=False, **kwargs):\n    \"\"\" MNASNet B1, depth multiplier of 1.0. \"\"\"\n    model = _gen_mnasnet_b1('mnasnet_100', 1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mnasnet_b1(pretrained=False, **kwargs):\n    \"\"\" MNASNet B1, depth multiplier of 1.0. \"\"\"\n    return mnasnet_100(pretrained, **kwargs)\n\n\ndef mnasnet_140(pretrained=False, **kwargs):\n    \"\"\" MNASNet B1,  depth multiplier of 1.4 \"\"\"\n    model = _gen_mnasnet_b1('mnasnet_140', 1.4, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef semnasnet_050(pretrained=False, **kwargs):\n    \"\"\" MNASNet A1 (w/ SE), depth multiplier of 0.5 \"\"\"\n    model = _gen_mnasnet_a1('semnasnet_050', 0.5, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef semnasnet_075(pretrained=False, **kwargs):\n    \"\"\" MNASNet A1 (w/ SE),  depth multiplier of 0.75. \"\"\"\n    model = _gen_mnasnet_a1('semnasnet_075', 0.75, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef semnasnet_100(pretrained=False, **kwargs):\n    \"\"\" MNASNet A1 (w/ SE), depth multiplier of 1.0. \"\"\"\n    model = _gen_mnasnet_a1('semnasnet_100', 1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mnasnet_a1(pretrained=False, **kwargs):\n    \"\"\" MNASNet A1 (w/ SE), depth multiplier of 1.0. \"\"\"\n    return semnasnet_100(pretrained, **kwargs)\n\n\ndef semnasnet_140(pretrained=False, **kwargs):\n    \"\"\" MNASNet A1 (w/ SE), depth multiplier of 1.4. \"\"\"\n    model = _gen_mnasnet_a1('semnasnet_140', 1.4, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mnasnet_small(pretrained=False, **kwargs):\n    \"\"\" MNASNet Small,  depth multiplier of 1.0. \"\"\"\n    model = _gen_mnasnet_small('mnasnet_small', 1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef fbnetc_100(pretrained=False, **kwargs):\n    \"\"\" FBNet-C \"\"\"\n    if pretrained:\n        # pretrained model trained with non-default BN epsilon\n        kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    model = _gen_fbnetc('fbnetc_100', 1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef spnasnet_100(pretrained=False, **kwargs):\n    \"\"\" Single-Path NAS Pixel1\"\"\"\n    model = _gen_spnasnet('spnasnet_100', 1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_b0(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B0 \"\"\"\n    # NOTE for train set drop_rate=0.2, drop_connect_rate=0.2\n    model = _gen_efficientnet(\n        'efficientnet_b0', channel_multiplier=1.0, depth_multiplier=1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_b1(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B1 \"\"\"\n    # NOTE for train set drop_rate=0.2, drop_connect_rate=0.2\n    model = _gen_efficientnet(\n        'efficientnet_b1', channel_multiplier=1.0, depth_multiplier=1.1, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_b2(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B2 \"\"\"\n    # NOTE for train set drop_rate=0.3, drop_connect_rate=0.2\n    model = _gen_efficientnet(\n        'efficientnet_b2', channel_multiplier=1.1, depth_multiplier=1.2, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_b3(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B3 \"\"\"\n    # NOTE for train set drop_rate=0.3, drop_connect_rate=0.2\n    model = _gen_efficientnet(\n        'efficientnet_b3', channel_multiplier=1.2, depth_multiplier=1.4, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_b4(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B4 \"\"\"\n    # NOTE for train set drop_rate=0.4, drop_connect_rate=0.2\n    model = _gen_efficientnet(\n        'efficientnet_b4', channel_multiplier=1.4, depth_multiplier=1.8, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_b5(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B5 \"\"\"\n    # NOTE for train set drop_rate=0.4, drop_connect_rate=0.2\n    model = _gen_efficientnet(\n        'efficientnet_b5', channel_multiplier=1.6, depth_multiplier=2.2, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_b6(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B6 \"\"\"\n    # NOTE for train set drop_rate=0.5, drop_connect_rate=0.2\n    model = _gen_efficientnet(\n        'efficientnet_b6', channel_multiplier=1.8, depth_multiplier=2.6, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_b7(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B7 \"\"\"\n    # NOTE for train set drop_rate=0.5, drop_connect_rate=0.2\n    model = _gen_efficientnet(\n        'efficientnet_b7', channel_multiplier=2.0, depth_multiplier=3.1, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_b8(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B8 \"\"\"\n    # NOTE for train set drop_rate=0.5, drop_connect_rate=0.2\n    model = _gen_efficientnet(\n        'efficientnet_b8', channel_multiplier=2.2, depth_multiplier=3.6, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_es(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-Edge Small. \"\"\"\n    model = _gen_efficientnet_edge(\n        'efficientnet_es', channel_multiplier=1.0, depth_multiplier=1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_em(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-Edge-Medium. \"\"\"\n    model = _gen_efficientnet_edge(\n        'efficientnet_em', channel_multiplier=1.0, depth_multiplier=1.1, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_el(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-Edge-Large. \"\"\"\n    model = _gen_efficientnet_edge(\n        'efficientnet_el', channel_multiplier=1.2, depth_multiplier=1.4, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_cc_b0_4e(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-CondConv-B0 w/ 8 Experts \"\"\"\n    # NOTE for train set drop_rate=0.25, drop_connect_rate=0.2\n    model = _gen_efficientnet_condconv(\n        'efficientnet_cc_b0_4e', channel_multiplier=1.0, depth_multiplier=1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_cc_b0_8e(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-CondConv-B0 w/ 8 Experts \"\"\"\n    # NOTE for train set drop_rate=0.25, drop_connect_rate=0.2\n    model = _gen_efficientnet_condconv(\n        'efficientnet_cc_b0_8e', channel_multiplier=1.0, depth_multiplier=1.0, experts_multiplier=2,\n        pretrained=pretrained, **kwargs)\n    return model\n\n\ndef efficientnet_cc_b1_8e(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-CondConv-B1 w/ 8 Experts \"\"\"\n    # NOTE for train set drop_rate=0.25, drop_connect_rate=0.2\n    model = _gen_efficientnet_condconv(\n        'efficientnet_cc_b1_8e', channel_multiplier=1.0, depth_multiplier=1.1, experts_multiplier=2,\n        pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b0(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B0. Tensorflow compatible variant  \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b0', channel_multiplier=1.0, depth_multiplier=1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b1(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B1. Tensorflow compatible variant  \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b1', channel_multiplier=1.0, depth_multiplier=1.1, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b2(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B2. Tensorflow compatible variant  \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b2', channel_multiplier=1.1, depth_multiplier=1.2, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b3(pretrained=False, num_classes=1000, in_chans=3, **kwargs):\n    \"\"\" EfficientNet-B3. Tensorflow compatible variant \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b3', channel_multiplier=1.2, depth_multiplier=1.4, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b4(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B4. Tensorflow compatible variant \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b4', channel_multiplier=1.4, depth_multiplier=1.8, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b5(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B5. Tensorflow compatible variant \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b5', channel_multiplier=1.6, depth_multiplier=2.2, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b6(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B6. Tensorflow compatible variant \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b6', channel_multiplier=1.8, depth_multiplier=2.6, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b7(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B7. Tensorflow compatible variant \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b7', channel_multiplier=2.0, depth_multiplier=3.1, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b0_ap(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B0 AdvProp. Tensorflow compatible variant  \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b0_ap', channel_multiplier=1.0, depth_multiplier=1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b1_ap(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B1 AdvProp. Tensorflow compatible variant  \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b1_ap', channel_multiplier=1.0, depth_multiplier=1.1, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b2_ap(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B2 AdvProp. Tensorflow compatible variant  \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b2_ap', channel_multiplier=1.1, depth_multiplier=1.2, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b3_ap(pretrained=False, num_classes=1000, in_chans=3, **kwargs):\n    \"\"\" EfficientNet-B3 AdvProp. Tensorflow compatible variant \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b3_ap', channel_multiplier=1.2, depth_multiplier=1.4, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b4_ap(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B4 AdvProp. Tensorflow compatible variant \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b4_ap', channel_multiplier=1.4, depth_multiplier=1.8, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b5_ap(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B5 AdvProp. Tensorflow compatible variant \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b5_ap', channel_multiplier=1.6, depth_multiplier=2.2, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b6_ap(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B6 AdvProp. Tensorflow compatible variant \"\"\"\n    # NOTE for train, drop_rate should be 0.5\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b6_ap', channel_multiplier=1.8, depth_multiplier=2.6, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b7_ap(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B7 AdvProp. Tensorflow compatible variant \"\"\"\n    # NOTE for train, drop_rate should be 0.5\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b7_ap', channel_multiplier=2.0, depth_multiplier=3.1, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_b8_ap(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-B8 AdvProp. Tensorflow compatible variant \"\"\"\n    # NOTE for train, drop_rate should be 0.5\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet(\n        'tf_efficientnet_b8_ap', channel_multiplier=2.2, depth_multiplier=3.6, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_es(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-Edge Small. Tensorflow compatible variant  \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet_edge(\n        'tf_efficientnet_es', channel_multiplier=1.0, depth_multiplier=1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_em(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-Edge-Medium. Tensorflow compatible variant  \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet_edge(\n        'tf_efficientnet_em', channel_multiplier=1.0, depth_multiplier=1.1, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_el(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-Edge-Large. Tensorflow compatible variant  \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet_edge(\n        'tf_efficientnet_el', channel_multiplier=1.2, depth_multiplier=1.4, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_cc_b0_4e(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-CondConv-B0 w/ 4 Experts \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet_condconv(\n        'tf_efficientnet_cc_b0_4e', channel_multiplier=1.0, depth_multiplier=1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\n\ndef tf_efficientnet_cc_b0_8e(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-CondConv-B0 w/ 8 Experts \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet_condconv(\n        'tf_efficientnet_cc_b0_8e', channel_multiplier=1.0, depth_multiplier=1.0, experts_multiplier=2,\n        pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_efficientnet_cc_b1_8e(pretrained=False, **kwargs):\n    \"\"\" EfficientNet-CondConv-B1 w/ 8 Experts \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_efficientnet_condconv(\n        'tf_efficientnet_cc_b1_8e', channel_multiplier=1.0, depth_multiplier=1.1, experts_multiplier=2,\n        pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mixnet_s(pretrained=False, **kwargs):\n    \"\"\"Creates a MixNet Small model.\n    \"\"\"\n    # NOTE for train set drop_rate=0.2\n    model = _gen_mixnet_s(\n        'mixnet_s', channel_multiplier=1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mixnet_m(pretrained=False, **kwargs):\n    \"\"\"Creates a MixNet Medium model.\n    \"\"\"\n    # NOTE for train set drop_rate=0.25\n    model = _gen_mixnet_m(\n        'mixnet_m', channel_multiplier=1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mixnet_l(pretrained=False, **kwargs):\n    \"\"\"Creates a MixNet Large model.\n    \"\"\"\n    # NOTE for train set drop_rate=0.25\n    model = _gen_mixnet_m(\n        'mixnet_l', channel_multiplier=1.3, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mixnet_xl(pretrained=False, **kwargs):\n    \"\"\"Creates a MixNet Extra-Large model.\n    Not a paper spec, experimental def by RW w/ depth scaling.\n    \"\"\"\n    # NOTE for train set drop_rate=0.25, drop_connect_rate=0.2\n    model = _gen_mixnet_m(\n        'mixnet_xl', channel_multiplier=1.6, depth_multiplier=1.2, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mixnet_xxl(pretrained=False, **kwargs):\n    \"\"\"Creates a MixNet Double Extra Large model.\n    Not a paper spec, experimental def by RW w/ depth scaling.\n    \"\"\"\n    # NOTE for train set drop_rate=0.3, drop_connect_rate=0.2\n    model = _gen_mixnet_m(\n        'mixnet_xxl', channel_multiplier=2.4, depth_multiplier=1.3, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_mixnet_s(pretrained=False, **kwargs):\n    \"\"\"Creates a MixNet Small model. Tensorflow compatible variant\n    \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_mixnet_s(\n        'tf_mixnet_s', channel_multiplier=1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_mixnet_m(pretrained=False, **kwargs):\n    \"\"\"Creates a MixNet Medium model. Tensorflow compatible variant\n    \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_mixnet_m(\n        'tf_mixnet_m', channel_multiplier=1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_mixnet_l(pretrained=False, **kwargs):\n    \"\"\"Creates a MixNet Large model. Tensorflow compatible variant\n    \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_mixnet_m(\n        'tf_mixnet_l', channel_multiplier=1.3, pretrained=pretrained, **kwargs)\n    return model\n"
  },
  {
    "path": "backbones/geffnet/helpers.py",
    "content": "import torch\nimport os\nfrom collections import OrderedDict\ntry:\n    from torch.hub import load_state_dict_from_url\nexcept ImportError:\n    from torch.utils.model_zoo import load_url as load_state_dict_from_url\n\n\ndef load_checkpoint(model, checkpoint_path):\n    if checkpoint_path and os.path.isfile(checkpoint_path):\n        print(\"=> Loading checkpoint '{}'\".format(checkpoint_path))\n        checkpoint = torch.load(checkpoint_path)\n        if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:\n            new_state_dict = OrderedDict()\n            for k, v in checkpoint['state_dict'].items():\n                if k.startswith('module'):\n                    name = k[7:]  # remove `module.`\n                else:\n                    name = k\n                new_state_dict[name] = v\n            model.load_state_dict(new_state_dict)\n        else:\n            model.load_state_dict(checkpoint)\n        print(\"=> Loaded checkpoint '{}'\".format(checkpoint_path))\n    else:\n        print(\"=> Error: No checkpoint found at '{}'\".format(checkpoint_path))\n        raise FileNotFoundError()\n\n\ndef load_pretrained(model, url, filter_fn=None, strict=True):\n    state_dict = load_state_dict_from_url(url, progress=False, map_location='cpu')\n\n    input_conv = 'conv_stem'\n    classifier = 'classifier'\n    in_chans = getattr(model, input_conv).weight.shape[1]\n    num_classes = getattr(model, classifier).weight.shape[0]\n\n    input_conv_weight = input_conv + '.weight'\n    pretrained_in_chans = state_dict[input_conv_weight].shape[1]\n    if in_chans != pretrained_in_chans:\n        if in_chans == 1:\n            print('=> Converting pretrained input conv {} from {} to 1 channel'.format(\n                input_conv_weight, pretrained_in_chans))\n            conv1_weight = state_dict[input_conv_weight]\n            state_dict[input_conv_weight] = conv1_weight.sum(dim=1, keepdim=True)\n        else:\n            print('=> Discarding pretrained input conv {} since input channel count != {}'.format(\n                input_conv_weight, pretrained_in_chans))\n            del state_dict[input_conv_weight]\n            strict = False\n\n    classifier_weight = classifier + '.weight'\n    pretrained_num_classes = state_dict[classifier_weight].shape[0]\n    if num_classes != pretrained_num_classes:\n        print('=> Discarding pretrained classifier since num_classes != {}'.format(pretrained_num_classes))\n        del state_dict[classifier_weight]\n        del state_dict[classifier + '.bias']\n        strict = False\n\n    if filter_fn is not None:\n        state_dict = filter_fn(state_dict)\n\n    model.load_state_dict(state_dict, strict=strict)\n"
  },
  {
    "path": "backbones/geffnet/mobilenetv3.py",
    "content": "\"\"\" MobileNet-V3\n\nA PyTorch impl of MobileNet-V3, compatible with TF weights from official impl.\n\nPaper: Searching for MobileNetV3 - https://arxiv.org/abs/1905.02244\n\nHacked together by Ross Wightman\n\"\"\"\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom .helpers import load_pretrained\nfrom .efficientnet_builder import *\n\n__all__ = ['mobilenetv3_rw', 'mobilenetv3_large_075', 'mobilenetv3_large_100', 'mobilenetv3_large_minimal_100',\n           'mobilenetv3_small_075', 'mobilenetv3_small_100', 'mobilenetv3_small_minimal_100',\n           'tf_mobilenetv3_large_075', 'tf_mobilenetv3_large_100', 'tf_mobilenetv3_large_minimal_100',\n           'tf_mobilenetv3_small_075', 'tf_mobilenetv3_small_100', 'tf_mobilenetv3_small_minimal_100']\n\nmodel_urls = {\n    'mobilenetv3_rw':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mobilenetv3_100-35495452.pth',\n    'mobilenetv3_large_075': None,\n    'mobilenetv3_large_100': None,\n    'mobilenetv3_large_minimal_100': None,\n    'mobilenetv3_small_075': None,\n    'mobilenetv3_small_100': None,\n    'mobilenetv3_small_minimal_100': None,\n    'tf_mobilenetv3_large_075':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_075-150ee8b0.pth',\n    'tf_mobilenetv3_large_100':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_100-427764d5.pth',\n    'tf_mobilenetv3_large_minimal_100':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_minimal_100-8596ae28.pth',\n    'tf_mobilenetv3_small_075':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_075-da427f52.pth',\n    'tf_mobilenetv3_small_100':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_100-37f49e2b.pth',\n    'tf_mobilenetv3_small_minimal_100':\n        'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_minimal_100-922a7843.pth',\n}\n\n\nclass MobileNetV3(nn.Module):\n    \"\"\" MobileNet-V3\n\n    A this model utilizes the MobileNet-v3 specific 'efficient head', where global pooling is done before the\n    head convolution without a final batch-norm layer before the classifier.\n\n    Paper: https://arxiv.org/abs/1905.02244\n    \"\"\"\n\n    def __init__(self, block_args, num_classes=1000, in_chans=3, stem_size=16, num_features=1280, head_bias=True,\n                 channel_multiplier=1.0, pad_type='', act_layer=HardSwish, drop_rate=0., drop_connect_rate=0.,\n                 se_kwargs=None, norm_layer=nn.BatchNorm2d, norm_kwargs=None, weight_init='goog'):\n        super(MobileNetV3, self).__init__()\n        self.drop_rate = drop_rate\n\n        stem_size = round_channels(stem_size, channel_multiplier)\n        self.conv_stem = select_conv2d(in_chans, stem_size, 3, stride=2, padding=pad_type)\n        self.bn1 = nn.BatchNorm2d(stem_size, **norm_kwargs)\n        self.act1 = act_layer(inplace=True)\n        in_chs = stem_size\n\n        builder = EfficientNetBuilder(\n            channel_multiplier, pad_type=pad_type, act_layer=act_layer, se_kwargs=se_kwargs,\n            norm_layer=norm_layer, norm_kwargs=norm_kwargs, drop_connect_rate=drop_connect_rate)\n        self.blocks = nn.Sequential(*builder(in_chs, block_args))\n        in_chs = builder.in_chs\n\n        self.global_pool = nn.AdaptiveAvgPool2d(1)\n        self.conv_head = select_conv2d(in_chs, num_features, 1, padding=pad_type, bias=head_bias)\n        self.act2 = act_layer(inplace=True)\n        self.classifier = nn.Linear(num_features, num_classes)\n\n        for m in self.modules():\n            if weight_init == 'goog':\n                initialize_weight_goog(m)\n            else:\n                initialize_weight_default(m)\n\n    def as_sequential(self):\n        layers = [self.conv_stem, self.bn1, self.act1]\n        layers.extend(self.blocks)\n        layers.extend([\n            self.global_pool, self.conv_head, self.act2,\n            nn.Flatten(), nn.Dropout(self.drop_rate), self.classifier])\n        return nn.Sequential(*layers)\n\n    def features(self, x):\n        x = self.conv_stem(x)\n        x = self.bn1(x)\n        x = self.act1(x)\n        x = self.blocks(x)\n        x = self.global_pool(x)\n        x = self.conv_head(x)\n        x = self.act2(x)\n        return x\n\n    def forward(self, x):\n        x = self.features(x)\n        x = x.flatten(1)\n        if self.drop_rate > 0.:\n            x = F.dropout(x, p=self.drop_rate, training=self.training)\n        return self.classifier(x)\n\n\ndef _create_model(model_kwargs, variant, pretrained=False):\n    as_sequential = model_kwargs.pop('as_sequential', False)\n    model = MobileNetV3(**model_kwargs)\n    if pretrained and model_urls[variant]:\n        load_pretrained(model, model_urls[variant])\n    if as_sequential:\n        model = model.as_sequential()\n    return model\n\n\ndef _gen_mobilenet_v3_rw(variant, channel_multiplier=1.0, pretrained=False, **kwargs):\n    \"\"\"Creates a MobileNet-V3 model (RW variant).\n\n    Paper: https://arxiv.org/abs/1905.02244\n\n    This was my first attempt at reproducing the MobileNet-V3 from paper alone. It came close to the\n    eventual Tensorflow reference impl but has a few differences:\n    1. This model has no bias on the head convolution\n    2. This model forces no residual (noskip) on the first DWS block, this is different than MnasNet\n    3. This model always uses ReLU for the SE activation layer, other models in the family inherit their act layer\n       from their parent block\n    4. This model does not enforce divisible by 8 limitation on the SE reduction channel count\n\n    Overall the changes are fairly minor and result in a very small parameter count difference and no\n    top-1/5\n\n    Args:\n      channel_multiplier: multiplier to number of channels per layer.\n    \"\"\"\n    arch_def = [\n        # stage 0, 112x112 in\n        ['ds_r1_k3_s1_e1_c16_nre_noskip'],  # relu\n        # stage 1, 112x112 in\n        ['ir_r1_k3_s2_e4_c24_nre', 'ir_r1_k3_s1_e3_c24_nre'],  # relu\n        # stage 2, 56x56 in\n        ['ir_r3_k5_s2_e3_c40_se0.25_nre'],  # relu\n        # stage 3, 28x28 in\n        ['ir_r1_k3_s2_e6_c80', 'ir_r1_k3_s1_e2.5_c80', 'ir_r2_k3_s1_e2.3_c80'],  # hard-swish\n        # stage 4, 14x14in\n        ['ir_r2_k3_s1_e6_c112_se0.25'],  # hard-swish\n        # stage 5, 14x14in\n        ['ir_r3_k5_s2_e6_c160_se0.25'],  # hard-swish\n        # stage 6, 7x7 in\n        ['cn_r1_k1_s1_c960'],  # hard-swish\n    ]\n    model_kwargs = dict(\n        block_args=decode_arch_def(arch_def),\n        head_bias=False,  # one of my mistakes\n        channel_multiplier=channel_multiplier,\n        act_layer=resolve_act_layer(kwargs, 'hard_swish'),\n        se_kwargs=dict(gate_fn=get_act_fn('hard_sigmoid'), reduce_mid=True),\n        norm_kwargs=resolve_bn_args(kwargs),\n        **kwargs,\n    )\n    model = _create_model(model_kwargs, variant, pretrained)\n    return model\n\n\ndef _gen_mobilenet_v3(variant, channel_multiplier=1.0, pretrained=False, **kwargs):\n    \"\"\"Creates a MobileNet-V3 large/small/minimal models.\n\n    Ref impl: https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet_v3.py\n    Paper: https://arxiv.org/abs/1905.02244\n\n    Args:\n      channel_multiplier: multiplier to number of channels per layer.\n    \"\"\"\n    if 'small' in variant:\n        num_features = 1024\n        if 'minimal' in variant:\n            act_layer = 'relu'\n            arch_def = [\n                # stage 0, 112x112 in\n                ['ds_r1_k3_s2_e1_c16'],\n                # stage 1, 56x56 in\n                ['ir_r1_k3_s2_e4.5_c24', 'ir_r1_k3_s1_e3.67_c24'],\n                # stage 2, 28x28 in\n                ['ir_r1_k3_s2_e4_c40', 'ir_r2_k3_s1_e6_c40'],\n                # stage 3, 14x14 in\n                ['ir_r2_k3_s1_e3_c48'],\n                # stage 4, 14x14in\n                ['ir_r3_k3_s2_e6_c96'],\n                # stage 6, 7x7 in\n                ['cn_r1_k1_s1_c576'],\n            ]\n        else:\n            act_layer = 'hard_swish'\n            arch_def = [\n                # stage 0, 112x112 in\n                ['ds_r1_k3_s2_e1_c16_se0.25_nre'],  # relu\n                # stage 1, 56x56 in\n                ['ir_r1_k3_s2_e4.5_c24_nre', 'ir_r1_k3_s1_e3.67_c24_nre'],  # relu\n                # stage 2, 28x28 in\n                ['ir_r1_k5_s2_e4_c40_se0.25', 'ir_r2_k5_s1_e6_c40_se0.25'],  # hard-swish\n                # stage 3, 14x14 in\n                ['ir_r2_k5_s1_e3_c48_se0.25'],  # hard-swish\n                # stage 4, 14x14in\n                ['ir_r3_k5_s2_e6_c96_se0.25'],  # hard-swish\n                # stage 6, 7x7 in\n                ['cn_r1_k1_s1_c576'],  # hard-swish\n            ]\n    else:\n        num_features = 1280\n        if 'minimal' in variant:\n            act_layer = 'relu'\n            arch_def = [\n                # stage 0, 112x112 in\n                ['ds_r1_k3_s1_e1_c16'],\n                # stage 1, 112x112 in\n                ['ir_r1_k3_s2_e4_c24', 'ir_r1_k3_s1_e3_c24'],\n                # stage 2, 56x56 in\n                ['ir_r3_k3_s2_e3_c40'],\n                # stage 3, 28x28 in\n                ['ir_r1_k3_s2_e6_c80', 'ir_r1_k3_s1_e2.5_c80', 'ir_r2_k3_s1_e2.3_c80'],\n                # stage 4, 14x14in\n                ['ir_r2_k3_s1_e6_c112'],\n                # stage 5, 14x14in\n                ['ir_r3_k3_s2_e6_c160'],\n                # stage 6, 7x7 in\n                ['cn_r1_k1_s1_c960'],\n            ]\n        else:\n            act_layer = 'hard_swish'\n            arch_def = [\n                # stage 0, 112x112 in\n                ['ds_r1_k3_s1_e1_c16_nre'],  # relu\n                # stage 1, 112x112 in\n                ['ir_r1_k3_s2_e4_c24_nre', 'ir_r1_k3_s1_e3_c24_nre'],  # relu\n                # stage 2, 56x56 in\n                ['ir_r3_k5_s2_e3_c40_se0.25_nre'],  # relu\n                # stage 3, 28x28 in\n                ['ir_r1_k3_s2_e6_c80', 'ir_r1_k3_s1_e2.5_c80', 'ir_r2_k3_s1_e2.3_c80'],  # hard-swish\n                # stage 4, 14x14in\n                ['ir_r2_k3_s1_e6_c112_se0.25'],  # hard-swish\n                # stage 5, 14x14in\n                ['ir_r3_k5_s2_e6_c160_se0.25'],  # hard-swish\n                # stage 6, 7x7 in\n                ['cn_r1_k1_s1_c960'],  # hard-swish\n            ]\n    model_kwargs = dict(\n        block_args=decode_arch_def(arch_def),\n        num_features=num_features,\n        stem_size=16,\n        channel_multiplier=channel_multiplier,\n        act_layer=resolve_act_layer(kwargs, act_layer),\n        se_kwargs=dict(\n            act_layer=get_act_layer('relu'), gate_fn=get_act_fn('hard_sigmoid'), reduce_mid=True, divisor=8),\n        norm_kwargs=resolve_bn_args(kwargs),\n        **kwargs,\n    )\n    model = _create_model(model_kwargs, variant, pretrained)\n    return model\n\n\ndef mobilenetv3_rw(pretrained=False, **kwargs):\n    \"\"\" MobileNet-V3 RW\n    Attn: See note in gen function for this variant.\n    \"\"\"\n    # NOTE for train set drop_rate=0.2\n    if pretrained:\n        # pretrained model trained with non-default BN epsilon\n        kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    model = _gen_mobilenet_v3_rw('mobilenetv3_rw', 1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mobilenetv3_large_075(pretrained=False, **kwargs):\n    \"\"\" MobileNet V3 Large 0.75\"\"\"\n    # NOTE for train set drop_rate=0.2\n    model = _gen_mobilenet_v3('mobilenetv3_large_075', 0.75, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mobilenetv3_large_100(pretrained=False, **kwargs):\n    \"\"\" MobileNet V3 Large 1.0 \"\"\"\n    # NOTE for train set drop_rate=0.2\n    model = _gen_mobilenet_v3('mobilenetv3_large_100', 1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mobilenetv3_large_minimal_100(pretrained=False, **kwargs):\n    \"\"\" MobileNet V3 Large (Minimalistic) 1.0 \"\"\"\n    # NOTE for train set drop_rate=0.2\n    model = _gen_mobilenet_v3('mobilenetv3_large_minimal_100', 1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mobilenetv3_small_075(pretrained=False, **kwargs):\n    \"\"\" MobileNet V3 Small 0.75 \"\"\"\n    model = _gen_mobilenet_v3('mobilenetv3_small_075', 0.75, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mobilenetv3_small_100(pretrained=False, **kwargs):\n    \"\"\" MobileNet V3 Small 1.0 \"\"\"\n    model = _gen_mobilenet_v3('mobilenetv3_small_100', 1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef mobilenetv3_small_minimal_100(pretrained=False, **kwargs):\n    \"\"\" MobileNet V3 Small (Minimalistic) 1.0 \"\"\"\n    model = _gen_mobilenet_v3('mobilenetv3_small_minimal_100', 1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_mobilenetv3_large_075(pretrained=False, **kwargs):\n    \"\"\" MobileNet V3 Large 0.75. Tensorflow compat variant. \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_mobilenet_v3('tf_mobilenetv3_large_075', 0.75, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_mobilenetv3_large_100(pretrained=False, **kwargs):\n    \"\"\" MobileNet V3 Large 1.0. Tensorflow compat variant. \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_mobilenet_v3('tf_mobilenetv3_large_100', 1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_mobilenetv3_large_minimal_100(pretrained=False, **kwargs):\n    \"\"\" MobileNet V3 Large Minimalistic 1.0. Tensorflow compat variant. \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_mobilenet_v3('tf_mobilenetv3_large_minimal_100', 1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_mobilenetv3_small_075(pretrained=False, **kwargs):\n    \"\"\" MobileNet V3 Small 0.75. Tensorflow compat variant. \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_mobilenet_v3('tf_mobilenetv3_small_075', 0.75, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_mobilenetv3_small_100(pretrained=False, **kwargs):\n    \"\"\" MobileNet V3 Small 1.0. Tensorflow compat variant.\"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_mobilenet_v3('tf_mobilenetv3_small_100', 1.0, pretrained=pretrained, **kwargs)\n    return model\n\n\ndef tf_mobilenetv3_small_minimal_100(pretrained=False, **kwargs):\n    \"\"\" MobileNet V3 Small Minimalistic 1.0. Tensorflow compat variant. \"\"\"\n    kwargs['bn_eps'] = BN_EPS_TF_DEFAULT\n    kwargs['pad_type'] = 'same'\n    model = _gen_mobilenet_v3('tf_mobilenetv3_small_minimal_100', 1.0, pretrained=pretrained, **kwargs)\n    return model\n"
  },
  {
    "path": "backbones/geffnet/model_factory.py",
    "content": "from .mobilenetv3 import *\nfrom .gen_efficientnet import *\nfrom .helpers import load_checkpoint\n\n\ndef create_model(\n        model_name='mnasnet_100',\n        pretrained=None,\n        num_classes=1000,\n        in_chans=3,\n        checkpoint_path='',\n        **kwargs):\n\n    margs = dict(num_classes=num_classes, in_chans=in_chans, pretrained=pretrained)\n\n    if model_name in globals():\n        create_fn = globals()[model_name]\n        model = create_fn(**margs, **kwargs)\n    else:\n        raise RuntimeError('Unknown model (%s)' % model_name)\n\n    if checkpoint_path and not pretrained:\n        load_checkpoint(model, checkpoint_path)\n\n    return model\n"
  },
  {
    "path": "backbones/geffnet/version.py",
    "content": "__version__ = '0.9.5'\n"
  },
  {
    "path": "configs/efficientdet_d2_bifpn_1x.py",
    "content": "# model settings\nnorm_cfg = dict(type='BN', requires_grad=False)\nmodel = dict(\n    type='RetinaNet',\n    backbone=dict(\n        type='EfficientNet',\n        model_name='tf_efficientnet_b2'),\n    neck=dict(\n        type='BIFPN',\n        in_channels=[48, 88, 120, 208, 352],\n        out_channels=112,\n        start_level=0,\n        stack=4,\n        add_extra_convs=True,\n        num_outs=5,\n        norm_cfg=dict(type='BN', requires_grad=False),\n        activation='relu'),\n    bbox_head=dict(\n        type='RetinaHead',\n        num_classes=81,\n        in_channels=112,#256->112\n        stacked_convs=3, #4->3\n        feat_channels=112,#256->112\n        octave_base_scale=4,\n        scales_per_octave=3,\n        anchor_ratios=[0.5, 1.0, 2.0],\n        anchor_strides=[8, 16, 32, 64, 128],\n        target_means=[.0, .0, .0, .0],\n        target_stds=[1.0, 1.0, 1.0, 1.0],\n        loss_cls=dict(\n            type='FocalLoss',\n            use_sigmoid=True,\n            gamma=1.5, #2->1.5\n            alpha=0.25,\n            loss_weight=1.0),\n        loss_bbox=dict(type='SmoothL1Loss', beta=0.11, loss_weight=1.0)))\n# training and testing settings\ntrain_cfg = dict(\n    assigner=dict(\n        type='MaxIoUAssigner',\n        pos_iou_thr=0.5,\n        neg_iou_thr=0.4,\n        min_pos_iou=0,\n        ignore_iof_thr=-1),\n    allowed_border=-1,\n    pos_weight=-1,\n    debug=False)\ntest_cfg = dict(\n    nms_pre=1000,\n    min_bbox_size=0,\n    score_thr=0.05,\n    nms=dict(type='nms', iou_thr=0.5),\n    max_per_img=100)\n# dataset settings\ndataset_type = 'CocoDataset'\ndata_root = 'data/coco/'\nimg_norm_cfg = dict(\n    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)\ntrain_pipeline = [\n    dict(type='LoadImageFromFile'),\n    dict(type='LoadAnnotations', with_bbox=True),\n    dict(type='Resize', img_scale=(768, 768), keep_ratio=True),\n    dict(type='RandomFlip', flip_ratio=0.5),\n    dict(type='Normalize', **img_norm_cfg),\n    dict(type='Pad', size_divisor=128),\n    dict(type='DefaultFormatBundle'),\n    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']),\n]\ntest_pipeline = [\n    dict(type='LoadImageFromFile'),\n    dict(\n        type='MultiScaleFlipAug',\n        img_scale=(768, 768),\n        flip=False,\n        transforms=[\n            dict(type='Resize', keep_ratio=True),\n            dict(type='RandomFlip'),\n            dict(type='Normalize', **img_norm_cfg),\n            dict(type='Pad', size_divisor=128),\n            dict(type='ImageToTensor', keys=['img']),\n            dict(type='Collect', keys=['img']),\n        ])\n]\ndata = dict(\n    imgs_per_gpu=2,\n    workers_per_gpu=2,\n    train=dict(\n        type=dataset_type,\n        ann_file=data_root + 'annotations/instances_train2017.json',\n        img_prefix=data_root + 'train2017/',\n        pipeline=train_pipeline),\n    val=dict(\n        type=dataset_type,\n        ann_file=data_root + 'annotations/instances_val2017.json',\n        img_prefix=data_root + 'val2017/',\n        pipeline=test_pipeline),\n    test=dict(\n        type=dataset_type,\n        ann_file=data_root + 'annotations/instances_val2017.json',\n        img_prefix=data_root + 'val2017/',\n        pipeline=test_pipeline))\n# optimizer\noptimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=4e-5) #wd 0.0001->4e-5\noptimizer_config = dict(grad_clip=dict(max_norm=35, norm_type=2))\n# learning policy\nlr_config = dict(\n    policy='step',\n    warmup='linear',\n    warmup_iters=500,\n    warmup_ratio=1.0 / 3,\n    step=[8, 11])\ncheckpoint_config = dict(interval=1)\n# yapf:disable\nlog_config = dict(\n    interval=50,\n    hooks=[\n        dict(type='TextLoggerHook'),\n        # dict(type='TensorboardLoggerHook')\n    ])\n# yapf:enable\n# runtime settings\ntotal_epochs = 12\ndist_params = dict(backend='nccl')\nlog_level = 'INFO'\nwork_dir = './work_dirs/efficient_d2_bifpn_1x'\nload_from = None\nresume_from = None\nworkflow = [('train', 1)]\n"
  },
  {
    "path": "configs/efficientdet_d4_bifpn_1x.py",
    "content": "# model settings\nmodel = dict(\n    type='RetinaNet',\n    backbone=dict(\n        type='EfficientNet',\n        model_name='tf_efficientnet_b4'),\n    neck=dict(\n        type='BIFPN',\n        in_channels=[56, 112, 160, 272, 448],\n        out_channels=224,\n        start_level=0,\n        stack=6,\n        add_extra_convs=True,\n        num_outs=5,\n        norm_cfg=dict(type='BN', requires_grad=False),\n        activation='relu'),\n    bbox_head=dict(\n        type='RetinaHead',\n        num_classes=81,\n        in_channels=224,#256->224\n        stacked_convs=4,\n        feat_channels=224,#256->224\n        octave_base_scale=4,\n        scales_per_octave=3,\n        anchor_ratios=[0.5, 1.0, 2.0],\n        anchor_strides=[8, 16, 32, 64, 128],\n        target_means=[.0, .0, .0, .0],\n        target_stds=[1.0, 1.0, 1.0, 1.0],\n        loss_cls=dict(\n            type='FocalLoss',\n            use_sigmoid=True,\n            gamma=1.5, #2->1.5\n            alpha=0.25,\n            loss_weight=1.0),\n        loss_bbox=dict(type='SmoothL1Loss', beta=0.11, loss_weight=1.0)))\n# training and testing settings\ntrain_cfg = dict(\n    assigner=dict(\n        type='MaxIoUAssigner',\n        pos_iou_thr=0.5,\n        neg_iou_thr=0.4,\n        min_pos_iou=0,\n        ignore_iof_thr=-1),\n    allowed_border=-1,\n    pos_weight=-1,\n    debug=False)\ntest_cfg = dict(\n    nms_pre=1000,\n    min_bbox_size=0,\n    score_thr=0.05,\n    nms=dict(type='nms', iou_thr=0.5),\n    max_per_img=100)\n# dataset settings\ndataset_type = 'CocoDataset'\ndata_root = 'data/coco/'\nimg_norm_cfg = dict(\n    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)\ntrain_pipeline = [\n    dict(type='LoadImageFromFile'),\n    dict(type='LoadAnnotations', with_bbox=True),\n    dict(type='Resize', img_scale=(1024, 1024), keep_ratio=True),\n    dict(type='RandomFlip', flip_ratio=0.5),\n    dict(type='Normalize', **img_norm_cfg),\n    dict(type='Pad', size_divisor=128),#32->128\n    dict(type='DefaultFormatBundle'),\n    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']),\n]\ntest_pipeline = [\n    dict(type='LoadImageFromFile'),\n    dict(\n        type='MultiScaleFlipAug',\n        img_scale=(1024, 1024),\n        flip=False,\n        transforms=[\n            dict(type='Resize', keep_ratio=True),\n            dict(type='RandomFlip'),\n            dict(type='Normalize', **img_norm_cfg),\n            dict(type='Pad', size_divisor=128),\n            dict(type='ImageToTensor', keys=['img']),\n            dict(type='Collect', keys=['img']),\n        ])\n]\ndata = dict(\n    imgs_per_gpu=2,\n    workers_per_gpu=2,\n    train=dict(\n        type=dataset_type,\n        ann_file=data_root + 'annotations/instances_train2017.json',\n        img_prefix=data_root + 'train2017/',\n        pipeline=train_pipeline),\n    val=dict(\n        type=dataset_type,\n        ann_file=data_root + 'annotations/instances_val2017.json',\n        img_prefix=data_root + 'val2017/',\n        pipeline=test_pipeline),\n    test=dict(\n        type=dataset_type,\n        ann_file=data_root + 'annotations/instances_val2017.json',\n        img_prefix=data_root + 'val2017/',\n        pipeline=test_pipeline))\n# optimizer\noptimizer = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=4e-5) #wd 0.0001->4e-5\noptimizer_config = dict(grad_clip=dict(max_norm=35, norm_type=2))\n# learning policy\nlr_config = dict(\n    policy='step',\n    warmup='linear',\n    warmup_iters=500,\n    warmup_ratio=1.0 / 3,\n    step=[8, 11])\ncheckpoint_config = dict(interval=1)\n# yapf:disable\nlog_config = dict(\n    interval=50,\n    hooks=[\n        dict(type='TextLoggerHook'),\n        # dict(type='TensorboardLoggerHook')\n    ])\n# yapf:enable\n# runtime settings\ntotal_epochs = 12\ndist_params = dict(backend='nccl')\nlog_level = 'INFO'\nwork_dir = './work_dirs/efficient_d4_bifpn_1x'\nload_from = None\nresume_from = None\nworkflow = [('train', 1)]\n"
  },
  {
    "path": "necks/bifpn.py",
    "content": "import torch.nn as nn\nimport torch.nn.functional as F\nfrom mmcv.cnn import xavier_init\n\nfrom mmdet.core import auto_fp16\nfrom ..registry import NECKS\nfrom ..utils import ConvModule\nimport torch\neps=0.0001\n\n@NECKS.register_module\nclass BIFPN(nn.Module):\n\n    def __init__(self,\n                 in_channels,\n                 out_channels,\n                 num_outs,\n                 start_level=0,\n                 end_level=-1,\n                 stack=1,\n                 add_extra_convs=False,\n                 extra_convs_on_inputs=True,\n                 relu_before_extra_convs=False,\n                 no_norm_on_lateral=False,\n                 conv_cfg=None,\n                 norm_cfg=None,\n                 activation=None):\n        super(BIFPN, self).__init__()\n        assert isinstance(in_channels, list)\n        self.in_channels = in_channels\n        self.out_channels = out_channels\n        self.num_ins = len(in_channels)\n        self.num_outs = num_outs\n        self.activation = activation\n        self.relu_before_extra_convs = relu_before_extra_convs\n        self.no_norm_on_lateral = no_norm_on_lateral\n        self.fp16_enabled = False\n        self.stack = stack\n\n        if end_level == -1:\n            self.backbone_end_level = self.num_ins\n            assert num_outs >= self.num_ins - start_level\n        else:\n            # if end_level < inputs, no extra level is allowed\n            self.backbone_end_level = end_level\n            assert end_level <= len(in_channels)\n            assert num_outs == end_level - start_level\n        self.start_level = start_level\n        self.end_level = end_level\n        self.add_extra_convs = add_extra_convs\n        self.extra_convs_on_inputs = extra_convs_on_inputs\n\n        self.lateral_convs = nn.ModuleList()\n        self.fpn_convs = nn.ModuleList()\n        self.stack_bifpn_convs = nn.ModuleList()\n\n        for i in range(self.start_level, self.backbone_end_level):\n            l_conv = ConvModule(\n                in_channels[i],\n                out_channels,\n                1,\n                conv_cfg=conv_cfg,\n                norm_cfg=norm_cfg if not self.no_norm_on_lateral else None,\n                activation=self.activation,\n                inplace=False)\n            self.lateral_convs.append(l_conv)\n\n        for ii in range(stack):\n            self.stack_bifpn_convs.append(BiFPNModule(channels=out_channels,\n                                                      levels= self.backbone_end_level-self.start_level,\n                                                      conv_cfg=conv_cfg,\n                                                      norm_cfg=norm_cfg,\n                                                      activation=activation))\n        # add extra conv layers (e.g., RetinaNet)\n        extra_levels = num_outs - self.backbone_end_level + self.start_level\n        if add_extra_convs and extra_levels >= 1:\n            for i in range(extra_levels):\n                if i == 0 and self.extra_convs_on_inputs:\n                    in_channels = self.in_channels[self.backbone_end_level - 1]\n                else:\n                    in_channels = out_channels\n                extra_fpn_conv = ConvModule(\n                    in_channels,\n                    out_channels,\n                    3,\n                    stride=2,\n                    padding=1,\n                    conv_cfg=conv_cfg,\n                    norm_cfg=norm_cfg,\n                    activation=self.activation,\n                    inplace=False)\n                self.fpn_convs.append(extra_fpn_conv)\n\n    # default init_weights for conv(msra) and norm in ConvModule\n    def init_weights(self):\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                xavier_init(m, distribution='uniform')\n\n    @auto_fp16()\n    def forward(self, inputs):\n        assert len(inputs) == len(self.in_channels)\n\n        # build laterals\n        laterals = [\n            lateral_conv(inputs[i + self.start_level])\n            for i, lateral_conv in enumerate(self.lateral_convs)\n        ]\n        # part 1: build top-down and down-top path with stack\n        used_backbone_levels = len(laterals)\n        for bifpn_module in self.stack_bifpn_convs:\n            laterals = bifpn_module(laterals)\n        outs = laterals\n        # part 2: add extra levels\n        if self.num_outs > len(outs):\n            # use max pool to get more levels on top of outputs\n            # (e.g., Faster R-CNN, Mask R-CNN)\n            if not self.add_extra_convs:\n                for i in range(self.num_outs - used_backbone_levels):\n                    outs.append(F.max_pool2d(outs[-1], 1, stride=2))\n            # add conv layers on top of original feature maps (RetinaNet)\n            else:\n                if self.extra_convs_on_inputs:\n                    orig = inputs[self.backbone_end_level - 1]\n                    outs.append(self.fpn_convs[0](orig))\n                else:\n                    outs.append(self.fpn_convs[0](outs[-1]))\n                for i in range(1, self.num_outs - used_backbone_levels):\n                    if self.relu_before_extra_convs:\n                        outs.append(self.fpn_convs[i](F.relu(outs[-1])))\n                    else:\n                        outs.append(self.fpn_convs[i](outs[-1]))\n        return tuple(outs)\n\n\nclass BiFPNModule(nn.Module):\n    def __init__(self,\n                 channels,\n                 levels,\n                 init=0.5,\n                 conv_cfg=None,\n                 norm_cfg=None,\n                 activation=None):\n        super(BiFPNModule, self).__init__()\n        self.activation = activation\n        self.levels = levels\n        self.bifpn_convs =nn.ModuleList()\n        #weighted\n        self.w1 = nn.Parameter(torch.Tensor(2, levels).fill_(init))\n        self.relu1 = nn.ReLU()\n        self.w2 = nn.Parameter(torch.Tensor(3, levels - 2).fill_(init))\n        self.relu2 = nn.ReLU()\n        for jj in range(2):\n            for i in range(self.levels-1):  # 1,2,3\n                fpn_conv = nn.Sequential(\n                    ConvModule(\n                        channels,\n                        channels,\n                        3,\n                        padding=1,\n                        groups=channels,\n                        conv_cfg=conv_cfg,\n                        norm_cfg=norm_cfg,\n                        activation=self.activation,\n                        inplace=False),\n                    ConvModule(\n                        channels,\n                        channels,\n                        1,\n                        conv_cfg=conv_cfg,\n                        norm_cfg=norm_cfg,\n                        activation=self.activation,\n                        inplace=False))\n                self.bifpn_convs.append(fpn_conv)\n\n    # default init_weights for conv(msra) and norm in ConvModule\n    def init_weights(self):\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                xavier_init(m, distribution='uniform')\n\n    @auto_fp16()\n    def forward(self, inputs):\n        assert len(inputs) == self.levels\n        # build top-down and down-top path with stack\n        levels = self.levels\n        #w relu\n        w1 = self.relu1(self.w1)\n        w1 /= torch.sum(w1, dim=0) + eps #normalize\n        w2 = self.relu2(self.w2)\n        w2 /= torch.sum(w2, dim=0) + eps\n        # build top-down\n        kk=0\n        pathtd = inputs\n        inputs_clone=[]\n        for in_tensor in inputs:\n            inputs_clone.append(in_tensor.clone())\n        for i in range(levels - 1, 0, -1):\n            pathtd[i - 1] = w1[0,kk]*pathtd[i - 1] + w1[1,kk]*F.interpolate(\n                pathtd[i], scale_factor=2, mode='nearest')\n            pathtd[i - 1] = self.bifpn_convs[kk](pathtd[i - 1])\n            kk=kk+1\n        jj=kk\n        # build down-top\n        for i in range(0, levels - 2, 1):\n            pathtd[i + 1] = w2[0, i] * pathtd[i + 1] + w2[1, i] * F.max_pool2d(pathtd[i], kernel_size=2) + w2[2, i] * \\\n                            inputs_clone[i + 1]\n            pathtd[i + 1] = self.bifpn_convs[jj](pathtd[i + 1])\n            jj=jj+1\n\n        pathtd[levels - 1] = w1[0, kk] * pathtd[levels - 1] + w1[1, kk] * F.max_pool2d(pathtd[levels - 2],\n                                                                                       kernel_size=2)\n        pathtd[levels - 1] = self.bifpn_convs[jj](pathtd[levels - 1])\n        return pathtd\n\n"
  }
]