From c2ba229d995c33aaaf20e00a5686b4dc857044be Mon Sep 17 00:00:00 2001
From: Ross Wightman <rwightman@gmail.com>
Date: Fri, 21 May 2021 17:47:49 -0700
Subject: [PATCH] Prep for effcientnetv2_rw_m model weights that started
 training before official release..

---
 timm/models/efficientnet.py         | 16 ++++++++++++++--
 timm/models/efficientnet_builder.py |  8 ++++++--
 2 files changed, 20 insertions(+), 4 deletions(-)

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