##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

## Created by: Hang Zhang

## Email: zhanghang0704@gmail.com

## Copyright (c) 2020

##

## LICENSE file in the root directory of this source tree 

##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

"""ResNeSt implemented in Gluon."""

__all__ = ['resnest50', 'resnest101',

           'resnest200', 'resnest269']

from .resnet import ResNet, Bottleneck

from mxnet import cpu

def resnest50(pretrained=False, root='~/.mxnet/models', ctx=cpu(0), **kwargs):

    model = ResNet(Bottleneck, [3, 4, 6, 3],

                      radix=2, cardinality=1, bottleneck_width=64,

                      deep_stem=True, avg_down=True,

                      avd=True, avd_first=False,

                      use_splat=True, dropblock_prob=0.1,

                      name_prefix='resnest_', **kwargs)

    if pretrained:

        from .model_store import get_model_file

        model.load_parameters(get_model_file('resnest50', root=root), ctx=ctx)

    return model

def resnest101(pretrained=False, root='~/.mxnet/models', ctx=cpu(0), **kwargs):

    model = ResNet(Bottleneck, [3, 4, 23, 3],

                      radix=2, cardinality=1, bottleneck_width=64,

                      deep_stem=True, avg_down=True, stem_width=64,

                      avd=True, avd_first=False, use_splat=True, dropblock_prob=0.1,

                      name_prefix='resnest_', **kwargs)

    if pretrained:

        from .model_store import get_model_file

        model.load_parameters(get_model_file('resnest101', root=root), ctx=ctx)

    return model

def resnest200(pretrained=False, root='~/.mxnet/models', ctx=cpu(0), **kwargs):

    model = ResNet(Bottleneck, [3, 24, 36, 3], deep_stem=True, avg_down=True, stem_width=64,

                      avd=True, use_splat=True, dropblock_prob=0.1, final_drop=0.2,

                      name_prefix='resnest_', **kwargs)

    if pretrained:

        from .model_store import get_model_file

        model.load_parameters(get_model_file('resnest200', root=root), ctx=ctx)

    return model

def resnest269(pretrained=False, root='~/.mxnet/models', ctx=cpu(0), **kwargs):

    model = ResNet(Bottleneck, [3, 30, 48, 8], deep_stem=True, avg_down=True, stem_width=64,

                      avd=True, use_splat=True, dropblock_prob=0.1, final_drop=0.2,

                      name_prefix='resnest_', **kwargs)

    if pretrained:

        from .model_store import get_model_file

        model.load_parameters(get_model_file('resnest269', root=root), ctx=ctx)

    return model