Skip to content

Commit 39e92f0

Browse files
committed
mobilenet_edgetpu can use group_size override, more consistency in arg wrap/sadface w/ extra group_size arg
1 parent b9f020a commit 39e92f0

File tree

2 files changed

+27
-19
lines changed

2 files changed

+27
-19
lines changed

timm/models/efficientnet.py

Lines changed: 24 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -488,7 +488,8 @@ def _gen_mnasnet_small(variant, channel_multiplier=1.0, pretrained=False, **kwar
488488

489489
def _gen_mobilenet_v1(
490490
variant, channel_multiplier=1.0, depth_multiplier=1.0,
491-
group_size=None, fix_stem_head=False, head_conv=False, pretrained=False, **kwargs):
491+
group_size=None, fix_stem_head=False, head_conv=False, pretrained=False, **kwargs
492+
):
492493
"""
493494
Ref impl: https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet_v2.py
494495
Paper: https://arxiv.org/abs/1801.04381
@@ -625,7 +626,8 @@ def _gen_spnasnet(variant, channel_multiplier=1.0, pretrained=False, **kwargs):
625626

626627
def _gen_efficientnet(
627628
variant, channel_multiplier=1.0, depth_multiplier=1.0, channel_divisor=8,
628-
group_size=None, pretrained=False, **kwargs):
629+
group_size=None, pretrained=False, **kwargs
630+
):
629631
"""Creates an EfficientNet model.
630632
631633
Ref impl: https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/efficientnet_model.py
@@ -673,7 +675,8 @@ def _gen_efficientnet(
673675

674676

675677
def _gen_efficientnet_edge(
676-
variant, channel_multiplier=1.0, depth_multiplier=1.0, group_size=None, pretrained=False, **kwargs):
678+
variant, channel_multiplier=1.0, depth_multiplier=1.0, group_size=None, pretrained=False, **kwargs
679+
):
677680
""" Creates an EfficientNet-EdgeTPU model
678681
679682
Ref impl: https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/edgetpu
@@ -704,7 +707,8 @@ def _gen_efficientnet_edge(
704707

705708

706709
def _gen_efficientnet_condconv(
707-
variant, channel_multiplier=1.0, depth_multiplier=1.0, experts_multiplier=1, pretrained=False, **kwargs):
710+
variant, channel_multiplier=1.0, depth_multiplier=1.0, experts_multiplier=1, pretrained=False, **kwargs
711+
):
708712
"""Creates an EfficientNet-CondConv model.
709713
710714
Ref impl: https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/condconv
@@ -776,7 +780,8 @@ def _gen_efficientnet_lite(variant, channel_multiplier=1.0, depth_multiplier=1.0
776780

777781

778782
def _gen_efficientnetv2_base(
779-
variant, channel_multiplier=1.0, depth_multiplier=1.0, group_size=None, pretrained=False, **kwargs):
783+
variant, channel_multiplier=1.0, depth_multiplier=1.0, group_size=None, pretrained=False, **kwargs
784+
):
780785
""" Creates an EfficientNet-V2 base model
781786
782787
Ref impl: https://github.com/google/automl/tree/master/efficientnetv2
@@ -805,7 +810,8 @@ def _gen_efficientnetv2_base(
805810

806811

807812
def _gen_efficientnetv2_s(
808-
variant, channel_multiplier=1.0, depth_multiplier=1.0, group_size=None, rw=False, pretrained=False, **kwargs):
813+
variant, channel_multiplier=1.0, depth_multiplier=1.0, group_size=None, rw=False, pretrained=False, **kwargs
814+
):
809815
""" Creates an EfficientNet-V2 Small model
810816
811817
Ref impl: https://github.com/google/automl/tree/master/efficientnetv2
@@ -844,7 +850,8 @@ def _gen_efficientnetv2_s(
844850

845851

846852
def _gen_efficientnetv2_m(
847-
variant, channel_multiplier=1.0, depth_multiplier=1.0, group_size=None, pretrained=False, **kwargs):
853+
variant, channel_multiplier=1.0, depth_multiplier=1.0, group_size=None, pretrained=False, **kwargs
854+
):
848855
""" Creates an EfficientNet-V2 Medium model
849856
850857
Ref impl: https://github.com/google/automl/tree/master/efficientnetv2
@@ -875,7 +882,8 @@ def _gen_efficientnetv2_m(
875882

876883

877884
def _gen_efficientnetv2_l(
878-
variant, channel_multiplier=1.0, depth_multiplier=1.0, group_size=None, pretrained=False, **kwargs):
885+
variant, channel_multiplier=1.0, depth_multiplier=1.0, group_size=None, pretrained=False, **kwargs
886+
):
879887
""" Creates an EfficientNet-V2 Large model
880888
881889
Ref impl: https://github.com/google/automl/tree/master/efficientnetv2
@@ -906,7 +914,8 @@ def _gen_efficientnetv2_l(
906914

907915

908916
def _gen_efficientnetv2_xl(
909-
variant, channel_multiplier=1.0, depth_multiplier=1.0, group_size=None, pretrained=False, **kwargs):
917+
variant, channel_multiplier=1.0, depth_multiplier=1.0, group_size=None, pretrained=False, **kwargs
918+
):
910919
""" Creates an EfficientNet-V2 Xtra-Large model
911920
912921
Ref impl: https://github.com/google/automl/tree/master/efficientnetv2
@@ -938,7 +947,8 @@ def _gen_efficientnetv2_xl(
938947

939948
def _gen_efficientnet_x(
940949
variant, channel_multiplier=1.0, depth_multiplier=1.0, channel_divisor=8,
941-
group_size=None, version=1, pretrained=False, **kwargs):
950+
group_size=None, version=1, pretrained=False, **kwargs
951+
):
942952
"""Creates an EfficientNet model.
943953
944954
Ref impl: https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/efficientnet_model.py
@@ -1084,9 +1094,7 @@ def _gen_mixnet_m(variant, channel_multiplier=1.0, depth_multiplier=1.0, pretrai
10841094
return model
10851095

10861096

1087-
def _gen_tinynet(
1088-
variant, model_width=1.0, depth_multiplier=1.0, pretrained=False, **kwargs
1089-
):
1097+
def _gen_tinynet(variant, model_width=1.0, depth_multiplier=1.0, pretrained=False, **kwargs):
10901098
"""Creates a TinyNet model.
10911099
"""
10921100
arch_def = [
@@ -1109,8 +1117,7 @@ def _gen_tinynet(
11091117
return model
11101118

11111119

1112-
def _gen_mobilenet_edgetpu(
1113-
variant, channel_multiplier=1.0, depth_multiplier=1.0, group_size=None, pretrained=False, **kwargs):
1120+
def _gen_mobilenet_edgetpu(variant, channel_multiplier=1.0, depth_multiplier=1.0, pretrained=False, **kwargs):
11141121
"""
11151122
Based on definitions in: https://github.com/tensorflow/models/tree/d2427a562f401c9af118e47af2f030a0a5599f55/official/projects/edgetpu/vision
11161123
"""
@@ -1186,7 +1193,7 @@ def _arch_def(chs: List[int], group_size: int):
11861193
]
11871194

11881195
model_kwargs = dict(
1189-
block_args=decode_arch_def(arch_def, depth_multiplier, group_size=group_size),
1196+
block_args=decode_arch_def(arch_def, depth_multiplier),
11901197
num_features=num_features,
11911198
stem_size=stem_size,
11921199
stem_kernel_size=stem_kernel_size,
@@ -1199,8 +1206,7 @@ def _arch_def(chs: List[int], group_size: int):
11991206
return model
12001207

12011208

1202-
def _gen_test_efficientnet(
1203-
variant, channel_multiplier=1.0, depth_multiplier=1.0, pretrained=False, **kwargs):
1209+
def _gen_test_efficientnet(variant, channel_multiplier=1.0, depth_multiplier=1.0, pretrained=False, **kwargs):
12041210
""" Minimal test EfficientNet generator.
12051211
"""
12061212
arch_def = [

timm/models/mobilenetv3.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -412,7 +412,9 @@ def _create_mnv3(variant: str, pretrained: bool = False, **kwargs) -> MobileNetV
412412
return model
413413

414414

415-
def _gen_mobilenet_v3_rw(variant: str, channel_multiplier: float = 1.0, pretrained: bool = False, **kwargs) -> MobileNetV3:
415+
def _gen_mobilenet_v3_rw(
416+
variant: str, channel_multiplier: float = 1.0, pretrained: bool = False, **kwargs
417+
) -> MobileNetV3:
416418
"""Creates a MobileNet-V3 model.
417419
418420
Ref impl: ?

0 commit comments

Comments
 (0)