|
7 | 7 |
|
8 | 8 | from timm.data import IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD
|
9 | 9 | from timm.layers import (
|
10 |
| - SelectAdaptivePool2d, Linear, LayerType, PadType, RmsNorm2d, ConvNormAct, create_conv2d, get_norm_act_layer, |
11 |
| - to_2tuple |
| 10 | + SelectAdaptivePool2d, Linear, LayerType, RmsNorm2d, ConvNormAct, create_conv2d, get_norm_act_layer, to_2tuple |
12 | 11 | )
|
13 | 12 | from ._builder import build_model_with_cfg
|
14 | 13 | from ._efficientnet_blocks import SqueezeExcite, UniversalInvertedResidual
|
15 | 14 | from ._efficientnet_builder import BlockArgs, EfficientNetBuilder, decode_arch_def, efficientnet_init_weights, \
|
16 |
| - round_channels, resolve_act_layer |
| 15 | + round_channels |
17 | 16 | from ._features import feature_take_indices
|
18 | 17 | from ._features_fx import register_notrace_module
|
19 |
| -from ._manipulate import checkpoint_seq, checkpoint |
| 18 | +from ._manipulate import checkpoint_seq |
20 | 19 | from ._registry import generate_default_cfgs, register_model
|
21 | 20 |
|
22 | 21 | __all__ = ['MobileNetV5', 'MobileNetV5Encoder']
|
23 | 22 |
|
| 23 | +_GELU = partial(nn.GELU, approximate='tanh') |
| 24 | + |
24 | 25 |
|
25 | 26 | @register_notrace_module
|
26 | 27 | class MobileNetV5MultiScaleFusionAdapter(nn.Module):
|
@@ -56,7 +57,7 @@ def __init__(
|
56 | 57 | self.layer_scale_init_value = layer_scale_init_value
|
57 | 58 | self.noskip = noskip
|
58 | 59 |
|
59 |
| - act_layer = act_layer or nn.GELU |
| 60 | + act_layer = act_layer or _GELU |
60 | 61 | norm_layer = norm_layer or RmsNorm2d
|
61 | 62 | self.ffn = UniversalInvertedResidual(
|
62 | 63 | in_chs=self.in_channels,
|
@@ -154,7 +155,7 @@ def __init__(
|
154 | 155 | global_pool: Type of pooling to use for global pooling features of the FC head.
|
155 | 156 | """
|
156 | 157 | super().__init__()
|
157 |
| - act_layer = act_layer or nn.GELU |
| 158 | + act_layer = act_layer or _GELU |
158 | 159 | norm_layer = norm_layer or RmsNorm2d
|
159 | 160 | norm_act_layer = get_norm_act_layer(norm_layer, act_layer)
|
160 | 161 | se_layer = se_layer or SqueezeExcite
|
@@ -411,7 +412,7 @@ def __init__(
|
411 | 412 | layer_scale_init_value: Optional[float] = None,
|
412 | 413 | ):
|
413 | 414 | super().__init__()
|
414 |
| - act_layer = act_layer or nn.GELU |
| 415 | + act_layer = act_layer or _GELU |
415 | 416 | norm_layer = norm_layer or RmsNorm2d
|
416 | 417 | se_layer = se_layer or SqueezeExcite
|
417 | 418 | self.num_classes = 0 # Exists to satisfy ._hub module APIs.
|
@@ -761,7 +762,7 @@ def _gen_mobilenet_v5(
|
761 | 762 | fix_stem=channel_multiplier < 1.0,
|
762 | 763 | round_chs_fn=partial(round_channels, multiplier=channel_multiplier),
|
763 | 764 | norm_layer=RmsNorm2d,
|
764 |
| - act_layer=nn.GELU, |
| 765 | + act_layer=_GELU, |
765 | 766 | layer_scale_init_value=1e-5,
|
766 | 767 | )
|
767 | 768 | model_kwargs = dict(model_kwargs, **kwargs)
|
|
0 commit comments