|
|
|
|
|
|
|
""" MobileNet V3
|
|
|
|
|
|
|
|
A PyTorch impl of MobileNet-V3, compatible with TF weights from official impl.
|
|
|
|
|
|
|
|
Paper: Searching for MobileNetV3 - https://arxiv.org/abs/1905.02244
|
|
|
|
|
|
|
|
Hacked together by / Copyright 2021 Ross Wightman
|
|
|
|
"""
|
|
|
|
from functools import partial
|
|
|
|
from typing import List
|
|
|
|
|
|
|
|
import torch
|
Monster commit, activation refactor, VoVNet, norm_act improvements, more
* refactor activations into basic PyTorch, jit scripted, and memory efficient custom auto
* implement hard-mish, better grad for hard-swish
* add initial VovNet V1/V2 impl, fix #151
* VovNet and DenseNet first models to use NormAct layers (support BatchNormAct2d, EvoNorm, InplaceIABN)
* Wrap IABN for any models that use it
* make more models torchscript compatible (DPN, PNasNet, Res2Net, SelecSLS) and add tests
4 years ago
|
|
|
import torch.nn as nn
|
|
|
|
import torch.nn.functional as F
|
|
|
|
|
|
|
|
from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD
|
|
|
|
from .efficientnet_blocks import SqueezeExcite
|
|
|
|
from .efficientnet_builder import EfficientNetBuilder, decode_arch_def, efficientnet_init_weights,\
|
|
|
|
round_channels, resolve_bn_args, resolve_act_layer, BN_EPS_TF_DEFAULT
|
|
|
|
from .features import FeatureInfo, FeatureHooks
|
|
|
|
from .helpers import build_model_with_cfg, default_cfg_for_features
|
|
|
|
from .layers import SelectAdaptivePool2d, Linear, create_conv2d, get_act_fn, hard_sigmoid
|
|
|
|
from .registry import register_model
|
|
|
|
|
|
|
|
__all__ = ['MobileNetV3', 'MobileNetV3Features']
|
|
|
|
|
|
|
|
|
|
|
|
def _cfg(url='', **kwargs):
|
|
|
|
return {
|
|
|
|
'url': url, 'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': (1, 1),
|
|
|
|
'crop_pct': 0.875, 'interpolation': 'bilinear',
|
|
|
|
'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
|
|
|
|
'first_conv': 'conv_stem', 'classifier': 'classifier',
|
|
|
|
**kwargs
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
default_cfgs = {
|
|
|
|
'mobilenetv3_large_075': _cfg(url=''),
|
|
|
|
'mobilenetv3_large_100': _cfg(
|
|
|
|
interpolation='bicubic',
|
|
|
|
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mobilenetv3_large_100_ra-f55367f5.pth'),
|
|
|
|
'mobilenetv3_large_100_miil': _cfg(
|
|
|
|
interpolation='bilinear', mean=(0, 0, 0), std=(1, 1, 1),
|
|
|
|
url='https://miil-public-eu.oss-eu-central-1.aliyuncs.com/model-zoo/ImageNet_21K_P/models/timm/mobilenetv3_large_100_1k_miil_78_0.pth'),
|
|
|
|
'mobilenetv3_large_100_miil_in21k': _cfg(
|
|
|
|
interpolation='bilinear', mean=(0, 0, 0), std=(1, 1, 1),
|
|
|
|
url='https://miil-public-eu.oss-eu-central-1.aliyuncs.com/model-zoo/ImageNet_21K_P/models/timm/mobilenetv3_large_100_in21k_miil.pth', num_classes=11221),
|
|
|
|
'mobilenetv3_small_075': _cfg(url=''),
|
|
|
|
'mobilenetv3_small_100': _cfg(url=''),
|
|
|
|
|
|
|
|
'mobilenetv3_rw': _cfg(
|
|
|
|
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/mobilenetv3_100-35495452.pth',
|
|
|
|
interpolation='bicubic'),
|
|
|
|
|
|
|
|
'tf_mobilenetv3_large_075': _cfg(
|
|
|
|
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_075-150ee8b0.pth',
|
|
|
|
mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD),
|
|
|
|
'tf_mobilenetv3_large_100': _cfg(
|
|
|
|
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_100-427764d5.pth',
|
|
|
|
mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD),
|
|
|
|
'tf_mobilenetv3_large_minimal_100': _cfg(
|
|
|
|
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_large_minimal_100-8596ae28.pth',
|
|
|
|
mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD),
|
|
|
|
'tf_mobilenetv3_small_075': _cfg(
|
|
|
|
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_075-da427f52.pth',
|
|
|
|
mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD),
|
|
|
|
'tf_mobilenetv3_small_100': _cfg(
|
|
|
|
url= 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_100-37f49e2b.pth',
|
|
|
|
mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD),
|
|
|
|
'tf_mobilenetv3_small_minimal_100': _cfg(
|
|
|
|
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_mobilenetv3_small_minimal_100-922a7843.pth',
|
|
|
|
mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD),
|
|
|
|
|
|
|
|
'fbnetv3_b': _cfg(),
|
|
|
|
'fbnetv3_d': _cfg(),
|
|
|
|
'fbnetv3_g': _cfg(),
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
class MobileNetV3(nn.Module):
|
|
|
|
""" MobiletNet-V3
|
|
|
|
|
|
|
|
Based on my EfficientNet implementation and building blocks, this model utilizes the MobileNet-v3 specific
|
|
|
|
'efficient head', where global pooling is done before the head convolution without a final batch-norm
|
|
|
|
layer before the classifier.
|
|
|
|
|
|
|
|
Paper: https://arxiv.org/abs/1905.02244
|
|
|
|
"""
|
|
|
|
|
|
|
|
def __init__(self, block_args, num_classes=1000, in_chans=3, stem_size=16, num_features=1280, head_bias=True,
|
|
|
|
pad_type='', act_layer=None, norm_layer=None, se_layer=None, se_from_exp=True,
|
|
|
|
round_chs_fn=round_channels, drop_rate=0., drop_path_rate=0., global_pool='avg'):
|
|
|
|
super(MobileNetV3, self).__init__()
|
|
|
|
act_layer = act_layer or nn.ReLU
|
|
|
|
norm_layer = norm_layer or nn.BatchNorm2d
|
|
|
|
se_layer = se_layer or SqueezeExcite
|
|
|
|
self.num_classes = num_classes
|
|
|
|
self.num_features = num_features
|
|
|
|
self.drop_rate = drop_rate
|
|
|
|
|
|
|
|
# Stem
|
|
|
|
stem_size = round_chs_fn(stem_size)
|
|
|
|
self.conv_stem = create_conv2d(in_chans, stem_size, 3, stride=2, padding=pad_type)
|
|
|
|
self.bn1 = norm_layer(stem_size)
|
|
|
|
self.act1 = act_layer(inplace=True)
|
|
|
|
|
|
|
|
# Middle stages (IR/ER/DS Blocks)
|
|
|
|
builder = EfficientNetBuilder(
|
|
|
|
output_stride=32, pad_type=pad_type, round_chs_fn=round_chs_fn, se_from_exp=se_from_exp,
|
|
|
|
act_layer=act_layer, norm_layer=norm_layer, se_layer=se_layer, drop_path_rate=drop_path_rate)
|
|
|
|
self.blocks = nn.Sequential(*builder(stem_size, block_args))
|
|
|
|
self.feature_info = builder.features
|
|
|
|
head_chs = builder.in_chs
|
|
|
|
|
|
|
|
# Head + Pooling
|
|
|
|
self.global_pool = SelectAdaptivePool2d(pool_type=global_pool)
|
|
|
|
num_pooled_chs = head_chs * self.global_pool.feat_mult()
|
|
|
|
self.conv_head = create_conv2d(num_pooled_chs, self.num_features, 1, padding=pad_type, bias=head_bias)
|
|
|
|
self.act2 = act_layer(inplace=True)
|
|
|
|
self.classifier = Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
|
|
|
|
|
|
|
|
efficientnet_init_weights(self)
|
|
|
|
|
|
|
|
def as_sequential(self):
|
|
|
|
layers = [self.conv_stem, self.bn1, self.act1]
|
|
|
|
layers.extend(self.blocks)
|
|
|
|
layers.extend([self.global_pool, self.conv_head, self.act2])
|
|
|
|
layers.extend([nn.Flatten(), nn.Dropout(self.drop_rate), self.classifier])
|
|
|
|
return nn.Sequential(*layers)
|
|
|
|
|
|
|
|
def get_classifier(self):
|
|
|
|
return self.classifier
|
|
|
|
|
|
|
|
def reset_classifier(self, num_classes, global_pool='avg'):
|
|
|
|
self.num_classes = num_classes
|
|
|
|
# cannot meaningfully change pooling of efficient head after creation
|
|
|
|
self.global_pool = SelectAdaptivePool2d(pool_type=global_pool)
|
|
|
|
self.classifier = Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
|
|
|
|
|
|
|
|
def forward_features(self, x):
|
|
|
|
x = self.conv_stem(x)
|
|
|
|
x = self.bn1(x)
|
|
|
|
x = self.act1(x)
|
|
|
|
x = self.blocks(x)
|
|
|
|
x = self.global_pool(x)
|
|
|
|
x = self.conv_head(x)
|
|
|
|
x = self.act2(x)
|
|
|
|
return x
|
|
|
|
|
|
|
|
def forward(self, x):
|
|
|
|
x = self.forward_features(x)
|
|
|
|
if not self.global_pool.is_identity():
|
|
|
|
x = x.flatten(1)
|
|
|
|
if self.drop_rate > 0.:
|
|
|
|
x = F.dropout(x, p=self.drop_rate, training=self.training)
|
|
|
|
return self.classifier(x)
|
|
|
|
|
|
|
|
|
|
|
|
class MobileNetV3Features(nn.Module):
|
|
|
|
""" MobileNetV3 Feature Extractor
|
|
|
|
|
|
|
|
A work-in-progress feature extraction module for MobileNet-V3 to use as a backbone for segmentation
|
|
|
|
and object detection models.
|
|
|
|
"""
|
|
|
|
|
|
|
|
def __init__(self, block_args, out_indices=(0, 1, 2, 3, 4), feature_location='bottleneck', in_chans=3,
|
|
|
|
stem_size=16, output_stride=32, pad_type='', round_chs_fn=round_channels, se_from_exp=True,
|
|
|
|
act_layer=None, norm_layer=None, se_layer=None, drop_rate=0., drop_path_rate=0.):
|
|
|
|
super(MobileNetV3Features, self).__init__()
|
|
|
|
act_layer = act_layer or nn.ReLU
|
|
|
|
norm_layer = norm_layer or nn.BatchNorm2d
|
|
|
|
se_layer = se_layer or SqueezeExcite
|
|
|
|
self.drop_rate = drop_rate
|
|
|
|
|
|
|
|
# Stem
|
|
|
|
stem_size = round_chs_fn(stem_size)
|
|
|
|
self.conv_stem = create_conv2d(in_chans, stem_size, 3, stride=2, padding=pad_type)
|
|
|
|
self.bn1 = norm_layer(stem_size)
|
|
|
|
self.act1 = act_layer(inplace=True)
|
|
|
|
|
|
|
|
# Middle stages (IR/ER/DS Blocks)
|
|
|
|
builder = EfficientNetBuilder(
|
|
|
|
output_stride=output_stride, pad_type=pad_type, round_chs_fn=round_chs_fn, se_from_exp=se_from_exp,
|
|
|
|
act_layer=act_layer, norm_layer=norm_layer, se_layer=se_layer,
|
|
|
|
drop_path_rate=drop_path_rate, feature_location=feature_location)
|
|
|
|
self.blocks = nn.Sequential(*builder(stem_size, block_args))
|
|
|
|
self.feature_info = FeatureInfo(builder.features, out_indices)
|
|
|
|
self._stage_out_idx = {v['stage']: i for i, v in enumerate(self.feature_info) if i in out_indices}
|
|
|
|
|
|
|
|
efficientnet_init_weights(self)
|
|
|
|
|
|
|
|
# Register feature extraction hooks with FeatureHooks helper
|
|
|
|
self.feature_hooks = None
|
|
|
|
if feature_location != 'bottleneck':
|
|
|
|
hooks = self.feature_info.get_dicts(keys=('module', 'hook_type'))
|
|
|
|
self.feature_hooks = FeatureHooks(hooks, self.named_modules())
|
|
|
|
|
|
|
|
def forward(self, x) -> List[torch.Tensor]:
|
|
|
|
x = self.conv_stem(x)
|
|
|
|
x = self.bn1(x)
|
|
|
|
x = self.act1(x)
|
|
|
|
if self.feature_hooks is None:
|
|
|
|
features = []
|
|
|
|
if 0 in self._stage_out_idx:
|
|
|
|
features.append(x) # add stem out
|
|
|
|
for i, b in enumerate(self.blocks):
|
|
|
|
x = b(x)
|
|
|
|
if i + 1 in self._stage_out_idx:
|
|
|
|
features.append(x)
|
|
|
|
return features
|
|
|
|
else:
|
|
|
|
self.blocks(x)
|
|
|
|
out = self.feature_hooks.get_output(x.device)
|
|
|
|
return list(out.values())
|
|
|
|
|
|
|
|
|
|
|
|
def _create_mnv3(variant, pretrained=False, **kwargs):
|
|
|
|
features_only = False
|
|
|
|
model_cls = MobileNetV3
|
|
|
|
kwargs_filter = None
|
|
|
|
if kwargs.pop('features_only', False):
|
|
|
|
features_only = True
|
|
|
|
kwargs_filter = ('num_classes', 'num_features', 'head_conv', 'head_bias', 'global_pool')
|
|
|
|
model_cls = MobileNetV3Features
|
|
|
|
model = build_model_with_cfg(
|
|
|
|
model_cls, variant, pretrained,
|
|
|
|
default_cfg=default_cfgs[variant],
|
|
|
|
pretrained_strict=not features_only,
|
|
|
|
kwargs_filter=kwargs_filter,
|
|
|
|
**kwargs)
|
|
|
|
if features_only:
|
|
|
|
model.default_cfg = default_cfg_for_features(model.default_cfg)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
def _gen_mobilenet_v3_rw(variant, channel_multiplier=1.0, pretrained=False, **kwargs):
|
|
|
|
"""Creates a MobileNet-V3 model.
|
|
|
|
|
|
|
|
Ref impl: ?
|
|
|
|
Paper: https://arxiv.org/abs/1905.02244
|
|
|
|
|
|
|
|
Args:
|
|
|
|
channel_multiplier: multiplier to number of channels per layer.
|
|
|
|
"""
|
|
|
|
arch_def = [
|
|
|
|
# stage 0, 112x112 in
|
|
|
|
['ds_r1_k3_s1_e1_c16_nre_noskip'], # relu
|
|
|
|
# stage 1, 112x112 in
|
|
|
|
['ir_r1_k3_s2_e4_c24_nre', 'ir_r1_k3_s1_e3_c24_nre'], # relu
|
|
|
|
# stage 2, 56x56 in
|
|
|
|
['ir_r3_k5_s2_e3_c40_se0.25_nre'], # relu
|
|
|
|
# stage 3, 28x28 in
|
|
|
|
['ir_r1_k3_s2_e6_c80', 'ir_r1_k3_s1_e2.5_c80', 'ir_r2_k3_s1_e2.3_c80'], # hard-swish
|
|
|
|
# stage 4, 14x14in
|
|
|
|
['ir_r2_k3_s1_e6_c112_se0.25'], # hard-swish
|
|
|
|
# stage 5, 14x14in
|
|
|
|
['ir_r3_k5_s2_e6_c160_se0.25'], # hard-swish
|
|
|
|
# stage 6, 7x7 in
|
|
|
|
['cn_r1_k1_s1_c960'], # hard-swish
|
|
|
|
]
|
|
|
|
model_kwargs = dict(
|
|
|
|
block_args=decode_arch_def(arch_def),
|
|
|
|
head_bias=False,
|
|
|
|
round_chs_fn=partial(round_channels, multiplier=channel_multiplier),
|
|
|
|
norm_layer=partial(nn.BatchNorm2d, **resolve_bn_args(kwargs)),
|
Monster commit, activation refactor, VoVNet, norm_act improvements, more
* refactor activations into basic PyTorch, jit scripted, and memory efficient custom auto
* implement hard-mish, better grad for hard-swish
* add initial VovNet V1/V2 impl, fix #151
* VovNet and DenseNet first models to use NormAct layers (support BatchNormAct2d, EvoNorm, InplaceIABN)
* Wrap IABN for any models that use it
* make more models torchscript compatible (DPN, PNasNet, Res2Net, SelecSLS) and add tests
4 years ago
|
|
|
act_layer=resolve_act_layer(kwargs, 'hard_swish'),
|
|
|
|
se_layer=partial(SqueezeExcite, gate_layer='hard_sigmoid'),
|
|
|
|
**kwargs,
|
|
|
|
)
|
|
|
|
model = _create_mnv3(variant, pretrained, **model_kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
def _gen_mobilenet_v3(variant, channel_multiplier=1.0, pretrained=False, **kwargs):
|
|
|
|
"""Creates a MobileNet-V3 model.
|
|
|
|
|
|
|
|
Ref impl: ?
|
|
|
|
Paper: https://arxiv.org/abs/1905.02244
|
|
|
|
|
|
|
|
Args:
|
|
|
|
channel_multiplier: multiplier to number of channels per layer.
|
|
|
|
"""
|
|
|
|
if 'small' in variant:
|
|
|
|
num_features = 1024
|
|
|
|
if 'minimal' in variant:
|
Monster commit, activation refactor, VoVNet, norm_act improvements, more
* refactor activations into basic PyTorch, jit scripted, and memory efficient custom auto
* implement hard-mish, better grad for hard-swish
* add initial VovNet V1/V2 impl, fix #151
* VovNet and DenseNet first models to use NormAct layers (support BatchNormAct2d, EvoNorm, InplaceIABN)
* Wrap IABN for any models that use it
* make more models torchscript compatible (DPN, PNasNet, Res2Net, SelecSLS) and add tests
4 years ago
|
|
|
act_layer = resolve_act_layer(kwargs, 'relu')
|
|
|
|
arch_def = [
|
|
|
|
# stage 0, 112x112 in
|
|
|
|
['ds_r1_k3_s2_e1_c16'],
|
|
|
|
# stage 1, 56x56 in
|
|
|
|
['ir_r1_k3_s2_e4.5_c24', 'ir_r1_k3_s1_e3.67_c24'],
|
|
|
|
# stage 2, 28x28 in
|
|
|
|
['ir_r1_k3_s2_e4_c40', 'ir_r2_k3_s1_e6_c40'],
|
|
|
|
# stage 3, 14x14 in
|
|
|
|
['ir_r2_k3_s1_e3_c48'],
|
|
|
|
# stage 4, 14x14in
|
|
|
|
['ir_r3_k3_s2_e6_c96'],
|
|
|
|
# stage 6, 7x7 in
|
|
|
|
['cn_r1_k1_s1_c576'],
|
|
|
|
]
|
|
|
|
else:
|
Monster commit, activation refactor, VoVNet, norm_act improvements, more
* refactor activations into basic PyTorch, jit scripted, and memory efficient custom auto
* implement hard-mish, better grad for hard-swish
* add initial VovNet V1/V2 impl, fix #151
* VovNet and DenseNet first models to use NormAct layers (support BatchNormAct2d, EvoNorm, InplaceIABN)
* Wrap IABN for any models that use it
* make more models torchscript compatible (DPN, PNasNet, Res2Net, SelecSLS) and add tests
4 years ago
|
|
|
act_layer = resolve_act_layer(kwargs, 'hard_swish')
|
|
|
|
arch_def = [
|
|
|
|
# stage 0, 112x112 in
|
|
|
|
['ds_r1_k3_s2_e1_c16_se0.25_nre'], # relu
|
|
|
|
# stage 1, 56x56 in
|
|
|
|
['ir_r1_k3_s2_e4.5_c24_nre', 'ir_r1_k3_s1_e3.67_c24_nre'], # relu
|
|
|
|
# stage 2, 28x28 in
|
|
|
|
['ir_r1_k5_s2_e4_c40_se0.25', 'ir_r2_k5_s1_e6_c40_se0.25'], # hard-swish
|
|
|
|
# stage 3, 14x14 in
|
|
|
|
['ir_r2_k5_s1_e3_c48_se0.25'], # hard-swish
|
|
|
|
# stage 4, 14x14in
|
|
|
|
['ir_r3_k5_s2_e6_c96_se0.25'], # hard-swish
|
|
|
|
# stage 6, 7x7 in
|
|
|
|
['cn_r1_k1_s1_c576'], # hard-swish
|
|
|
|
]
|
|
|
|
else:
|
|
|
|
num_features = 1280
|
|
|
|
if 'minimal' in variant:
|
Monster commit, activation refactor, VoVNet, norm_act improvements, more
* refactor activations into basic PyTorch, jit scripted, and memory efficient custom auto
* implement hard-mish, better grad for hard-swish
* add initial VovNet V1/V2 impl, fix #151
* VovNet and DenseNet first models to use NormAct layers (support BatchNormAct2d, EvoNorm, InplaceIABN)
* Wrap IABN for any models that use it
* make more models torchscript compatible (DPN, PNasNet, Res2Net, SelecSLS) and add tests
4 years ago
|
|
|
act_layer = resolve_act_layer(kwargs, 'relu')
|
|
|
|
arch_def = [
|
|
|
|
# stage 0, 112x112 in
|
|
|
|
['ds_r1_k3_s1_e1_c16'],
|
|
|
|
# stage 1, 112x112 in
|
|
|
|
['ir_r1_k3_s2_e4_c24', 'ir_r1_k3_s1_e3_c24'],
|
|
|
|
# stage 2, 56x56 in
|
|
|
|
['ir_r3_k3_s2_e3_c40'],
|
|
|
|
# stage 3, 28x28 in
|
|
|
|
['ir_r1_k3_s2_e6_c80', 'ir_r1_k3_s1_e2.5_c80', 'ir_r2_k3_s1_e2.3_c80'],
|
|
|
|
# stage 4, 14x14in
|
|
|
|
['ir_r2_k3_s1_e6_c112'],
|
|
|
|
# stage 5, 14x14in
|
|
|
|
['ir_r3_k3_s2_e6_c160'],
|
|
|
|
# stage 6, 7x7 in
|
|
|
|
['cn_r1_k1_s1_c960'],
|
|
|
|
]
|
|
|
|
else:
|
Monster commit, activation refactor, VoVNet, norm_act improvements, more
* refactor activations into basic PyTorch, jit scripted, and memory efficient custom auto
* implement hard-mish, better grad for hard-swish
* add initial VovNet V1/V2 impl, fix #151
* VovNet and DenseNet first models to use NormAct layers (support BatchNormAct2d, EvoNorm, InplaceIABN)
* Wrap IABN for any models that use it
* make more models torchscript compatible (DPN, PNasNet, Res2Net, SelecSLS) and add tests
4 years ago
|
|
|
act_layer = resolve_act_layer(kwargs, 'hard_swish')
|
|
|
|
arch_def = [
|
|
|
|
# stage 0, 112x112 in
|
|
|
|
['ds_r1_k3_s1_e1_c16_nre'], # relu
|
|
|
|
# stage 1, 112x112 in
|
|
|
|
['ir_r1_k3_s2_e4_c24_nre', 'ir_r1_k3_s1_e3_c24_nre'], # relu
|
|
|
|
# stage 2, 56x56 in
|
|
|
|
['ir_r3_k5_s2_e3_c40_se0.25_nre'], # relu
|
|
|
|
# stage 3, 28x28 in
|
|
|
|
['ir_r1_k3_s2_e6_c80', 'ir_r1_k3_s1_e2.5_c80', 'ir_r2_k3_s1_e2.3_c80'], # hard-swish
|
|
|
|
# stage 4, 14x14in
|
|
|
|
['ir_r2_k3_s1_e6_c112_se0.25'], # hard-swish
|
|
|
|
# stage 5, 14x14in
|
|
|
|
['ir_r3_k5_s2_e6_c160_se0.25'], # hard-swish
|
|
|
|
# stage 6, 7x7 in
|
|
|
|
['cn_r1_k1_s1_c960'], # hard-swish
|
|
|
|
]
|
|
|
|
se_layer = partial(SqueezeExcite, gate_layer='hard_sigmoid', force_act_layer=nn.ReLU, rd_round_fn=round_channels)
|
|
|
|
model_kwargs = dict(
|
|
|
|
block_args=decode_arch_def(arch_def),
|
|
|
|
num_features=num_features,
|
|
|
|
stem_size=16,
|
|
|
|
round_chs_fn=partial(round_channels, multiplier=channel_multiplier),
|
|
|
|
norm_layer=partial(nn.BatchNorm2d, **resolve_bn_args(kwargs)),
|
|
|
|
act_layer=act_layer,
|
|
|
|
se_layer=se_layer,
|
|
|
|
**kwargs,
|
|
|
|
)
|
|
|
|
model = _create_mnv3(variant, pretrained, **model_kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
def _gen_fbnetv3(variant, channel_multiplier=1.0, pretrained=False, **kwargs):
|
|
|
|
""" FBNetV3
|
|
|
|
Paper: `FBNetV3: Joint Architecture-Recipe Search using Predictor Pretraining`
|
|
|
|
- https://arxiv.org/abs/2006.02049
|
|
|
|
FIXME untested, this is a preliminary impl of some FBNet-V3 variants.
|
|
|
|
"""
|
|
|
|
vl = variant.split('_')[-1]
|
|
|
|
if vl in ('a', 'b'):
|
|
|
|
stem_size = 16
|
|
|
|
arch_def = [
|
|
|
|
['ds_r2_k3_s1_e1_c16'],
|
|
|
|
['ir_r1_k5_s2_e4_c24', 'ir_r3_k5_s1_e2_c24'],
|
|
|
|
['ir_r1_k5_s2_e5_c40_se0.25', 'ir_r4_k5_s1_e3_c40_se0.25'],
|
|
|
|
['ir_r1_k5_s2_e5_c72', 'ir_r4_k3_s1_e3_c72'],
|
|
|
|
['ir_r1_k3_s1_e5_c120_se0.25', 'ir_r5_k5_s1_e3_c120_se0.25'],
|
|
|
|
['ir_r1_k3_s2_e6_c184_se0.25', 'ir_r5_k5_s1_e4_c184_se0.25', 'ir_r1_k5_s1_e6_c224_se0.25'],
|
|
|
|
['cn_r1_k1_s1_c1344'],
|
|
|
|
]
|
|
|
|
elif vl == 'd':
|
|
|
|
stem_size = 24
|
|
|
|
arch_def = [
|
|
|
|
['ds_r2_k3_s1_e1_c16'],
|
|
|
|
['ir_r1_k3_s2_e5_c24', 'ir_r5_k3_s1_e2_c24'],
|
|
|
|
['ir_r1_k5_s2_e4_c40_se0.25', 'ir_r4_k3_s1_e3_c40_se0.25'],
|
|
|
|
['ir_r1_k3_s2_e5_c72', 'ir_r4_k3_s1_e3_c72'],
|
|
|
|
['ir_r1_k3_s1_e5_c128_se0.25', 'ir_r6_k5_s1_e3_c128_se0.25'],
|
|
|
|
['ir_r1_k3_s2_e6_c208_se0.25', 'ir_r5_k5_s1_e5_c208_se0.25', 'ir_r1_k5_s1_e6_c240_se0.25'],
|
|
|
|
['cn_r1_k1_s1_c1440'],
|
|
|
|
]
|
|
|
|
elif vl == 'g':
|
|
|
|
stem_size = 32
|
|
|
|
arch_def = [
|
|
|
|
['ds_r3_k3_s1_e1_c24'],
|
|
|
|
['ir_r1_k5_s2_e4_c40', 'ir_r4_k5_s1_e2_c40'],
|
|
|
|
['ir_r1_k5_s2_e4_c56_se0.25', 'ir_r4_k5_s1_e3_c56_se0.25'],
|
|
|
|
['ir_r1_k5_s2_e5_c104', 'ir_r4_k3_s1_e3_c104'],
|
|
|
|
['ir_r1_k3_s1_e5_c160_se0.25', 'ir_r8_k5_s1_e3_c160_se0.25'],
|
|
|
|
['ir_r1_k3_s2_e6_c264_se0.25', 'ir_r6_k5_s1_e5_c264_se0.25', 'ir_r2_k5_s1_e6_c288_se0.25'],
|
|
|
|
['cn_r1_k1_s1_c1728'],
|
|
|
|
]
|
|
|
|
else:
|
|
|
|
raise NotImplemented
|
|
|
|
round_chs_fn = partial(round_channels, multiplier=channel_multiplier, round_limit=0.95)
|
|
|
|
se_layer = partial(SqueezeExcite, gate_layer='hard_sigmoid', rd_round_fn=round_chs_fn)
|
|
|
|
act_layer = resolve_act_layer(kwargs, 'hard_swish')
|
|
|
|
model_kwargs = dict(
|
|
|
|
block_args=decode_arch_def(arch_def),
|
|
|
|
num_features=1984,
|
|
|
|
head_bias=False,
|
|
|
|
stem_size=stem_size,
|
|
|
|
round_chs_fn=round_chs_fn,
|
|
|
|
se_from_exp=False,
|
|
|
|
norm_layer=partial(nn.BatchNorm2d, **resolve_bn_args(kwargs)),
|
|
|
|
act_layer=act_layer,
|
|
|
|
se_layer=se_layer,
|
|
|
|
**kwargs,
|
|
|
|
)
|
|
|
|
model = _create_mnv3(variant, pretrained, **model_kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def mobilenetv3_large_075(pretrained=False, **kwargs):
|
|
|
|
""" MobileNet V3 """
|
|
|
|
model = _gen_mobilenet_v3('mobilenetv3_large_075', 0.75, pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def mobilenetv3_large_100(pretrained=False, **kwargs):
|
|
|
|
""" MobileNet V3 """
|
|
|
|
model = _gen_mobilenet_v3('mobilenetv3_large_100', 1.0, pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def mobilenetv3_large_100_miil(pretrained=False, **kwargs):
|
|
|
|
""" MobileNet V3
|
|
|
|
Weights taken from: https://github.com/Alibaba-MIIL/ImageNet21K
|
|
|
|
"""
|
|
|
|
model = _gen_mobilenet_v3('mobilenetv3_large_100_miil', 1.0, pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def mobilenetv3_large_100_miil_in21k(pretrained=False, **kwargs):
|
|
|
|
""" MobileNet V3, 21k pretraining
|
|
|
|
Weights taken from: https://github.com/Alibaba-MIIL/ImageNet21K
|
|
|
|
"""
|
|
|
|
model = _gen_mobilenet_v3('mobilenetv3_large_100_miil_in21k', 1.0, pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def mobilenetv3_small_075(pretrained=False, **kwargs):
|
|
|
|
""" MobileNet V3 """
|
|
|
|
model = _gen_mobilenet_v3('mobilenetv3_small_075', 0.75, pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def mobilenetv3_small_100(pretrained=False, **kwargs):
|
|
|
|
""" MobileNet V3 """
|
|
|
|
model = _gen_mobilenet_v3('mobilenetv3_small_100', 1.0, pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def mobilenetv3_rw(pretrained=False, **kwargs):
|
|
|
|
""" MobileNet V3 """
|
|
|
|
if pretrained:
|
|
|
|
# pretrained model trained with non-default BN epsilon
|
|
|
|
kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
|
|
|
|
model = _gen_mobilenet_v3_rw('mobilenetv3_rw', 1.0, pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def tf_mobilenetv3_large_075(pretrained=False, **kwargs):
|
|
|
|
""" MobileNet V3 """
|
|
|
|
kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
|
|
|
|
kwargs['pad_type'] = 'same'
|
|
|
|
model = _gen_mobilenet_v3('tf_mobilenetv3_large_075', 0.75, pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def tf_mobilenetv3_large_100(pretrained=False, **kwargs):
|
|
|
|
""" MobileNet V3 """
|
|
|
|
kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
|
|
|
|
kwargs['pad_type'] = 'same'
|
|
|
|
model = _gen_mobilenet_v3('tf_mobilenetv3_large_100', 1.0, pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def tf_mobilenetv3_large_minimal_100(pretrained=False, **kwargs):
|
|
|
|
""" MobileNet V3 """
|
|
|
|
kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
|
|
|
|
kwargs['pad_type'] = 'same'
|
|
|
|
model = _gen_mobilenet_v3('tf_mobilenetv3_large_minimal_100', 1.0, pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def tf_mobilenetv3_small_075(pretrained=False, **kwargs):
|
|
|
|
""" MobileNet V3 """
|
|
|
|
kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
|
|
|
|
kwargs['pad_type'] = 'same'
|
|
|
|
model = _gen_mobilenet_v3('tf_mobilenetv3_small_075', 0.75, pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def tf_mobilenetv3_small_100(pretrained=False, **kwargs):
|
|
|
|
""" MobileNet V3 """
|
|
|
|
kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
|
|
|
|
kwargs['pad_type'] = 'same'
|
|
|
|
model = _gen_mobilenet_v3('tf_mobilenetv3_small_100', 1.0, pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def tf_mobilenetv3_small_minimal_100(pretrained=False, **kwargs):
|
|
|
|
""" MobileNet V3 """
|
|
|
|
kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
|
|
|
|
kwargs['pad_type'] = 'same'
|
|
|
|
model = _gen_mobilenet_v3('tf_mobilenetv3_small_minimal_100', 1.0, pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def fbnetv3_b(pretrained=False, **kwargs):
|
|
|
|
""" FBNetV3-B """
|
|
|
|
model = _gen_fbnetv3('fbnetv3_b', pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def fbnetv3_d(pretrained=False, **kwargs):
|
|
|
|
""" FBNetV3-D """
|
|
|
|
model = _gen_fbnetv3('fbnetv3_d', pretrained=pretrained, **kwargs)
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
@register_model
|
|
|
|
def fbnetv3_g(pretrained=False, **kwargs):
|
|
|
|
""" FBNetV3-G """
|
|
|
|
model = _gen_fbnetv3('fbnetv3_g', pretrained=pretrained, **kwargs)
|
|
|
|
return model
|