diff --git a/timm/models/efficientnet.py b/timm/models/efficientnet.py index 0c0464b5..37c1c745 100644 --- a/timm/models/efficientnet.py +++ b/timm/models/efficientnet.py @@ -162,6 +162,9 @@ default_cfgs = { 'efficientnetv2_rw_s': _cfg( url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/efficientnet_v2s_ra2_288-a6477665.pth', input_size=(3, 288, 288), test_input_size=(3, 384, 384), pool_size=(9, 9), crop_pct=1.0), + 'efficientnetv2_rw_m': _cfg( + url='', + input_size=(3, 320, 320), test_input_size=(3, 416, 416), pool_size=(10, 10), crop_pct=1.0), 'efficientnetv2_s': _cfg( url='', @@ -173,7 +176,6 @@ default_cfgs = { url='', input_size=(3, 384, 384), test_input_size=(3, 480, 480), pool_size=(12, 12), crop_pct=1.0), - 'tf_efficientnet_b0': _cfg( url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/tf_efficientnet_b0_aa-827b6e33.pth', input_size=(3, 224, 224)), @@ -1461,7 +1463,7 @@ def efficientnet_b3_pruned(pretrained=False, **kwargs): @register_model def efficientnetv2_rw_s(pretrained=False, **kwargs): - """ EfficientNet-V2 Small. + """ EfficientNet-V2 Small RW variant. NOTE: This is my initial (pre official code release) w/ some differences. See efficientnetv2_s and tf_efficientnetv2_s for versions that match the official w/ PyTorch vs TF padding """ @@ -1469,6 +1471,16 @@ def efficientnetv2_rw_s(pretrained=False, **kwargs): return model +@register_model +def efficientnetv2_rw_m(pretrained=False, **kwargs): + """ EfficientNet-V2 Medium RW variant. + """ + model = _gen_efficientnetv2_s( + 'efficientnetv2_rw_m', channel_multiplier=1.2, depth_multiplier=(1.2,) * 4 + (1.6,) * 2, rw=True, + pretrained=pretrained, **kwargs) + return model + + @register_model def efficientnetv2_s(pretrained=False, **kwargs): """ EfficientNet-V2 Small. """ diff --git a/timm/models/efficientnet_builder.py b/timm/models/efficientnet_builder.py index 9d5853c7..30739454 100644 --- a/timm/models/efficientnet_builder.py +++ b/timm/models/efficientnet_builder.py @@ -237,7 +237,11 @@ def _scale_stage_depth(stack_args, repeats, depth_multiplier=1.0, depth_trunc='c def decode_arch_def(arch_def, depth_multiplier=1.0, depth_trunc='ceil', experts_multiplier=1, fix_first_last=False): arch_args = [] - for stack_idx, block_strings in enumerate(arch_def): + if isinstance(depth_multiplier, tuple): + assert len(depth_multiplier) == len(arch_def) + else: + depth_multiplier = (depth_multiplier,) * len(arch_def) + for stack_idx, (block_strings, multiplier) in enumerate(zip(arch_def, depth_multiplier)): assert isinstance(block_strings, list) stack_args = [] repeats = [] @@ -251,7 +255,7 @@ def decode_arch_def(arch_def, depth_multiplier=1.0, depth_trunc='ceil', experts_ if fix_first_last and (stack_idx == 0 or stack_idx == len(arch_def) - 1): arch_args.append(_scale_stage_depth(stack_args, repeats, 1.0, depth_trunc)) else: - arch_args.append(_scale_stage_depth(stack_args, repeats, depth_multiplier, depth_trunc)) + arch_args.append(_scale_stage_depth(stack_args, repeats, multiplier, depth_trunc)) return arch_args