Phaze-A: Add EfficientNetV2 Encoder

o2team · May 4, 2022 · 198c682 · 198c682
1 parent 0189029
commit 198c682
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 2 deletions.
diff --git a/plugins/train/model/phaze_a.py b/plugins/train/model/phaze_a.py
@@ -53,6 +53,20 @@
         keras_name="EfficientNetB6", no_amd=True, tf_min=2.3, scaling=(0, 255), default_size=528),
     efficientnet_b7=dict(
         keras_name="EfficientNetB7", no_amd=True, tf_min=2.3, scaling=(0, 255), default_size=600),
+    efficientnet_v2_b0=dict(
+        keras_name="EfficientNetV2B0", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=224),
+    efficientnet_v2_b1=dict(
+        keras_name="EfficientNetV2B1", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=240),
+    efficientnet_v2_b2=dict(
+        keras_name="EfficientNetV2B2", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=260),
+    efficientnet_v2_b3=dict(
+        keras_name="EfficientNetV2B3", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=300),
+    efficientnet_v2_s=dict(
+        keras_name="EfficientNetV2S", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=384),
+    efficientnet_v2_m=dict(
+        keras_name="EfficientNetV2M", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=480),
+    efficientnet_v2_l=dict(
+        keras_name="EfficientNetV2L", no_amd=True, tf_min=2.8, scaling=(-1, 1), default_size=480),
     inception_resnet_v2=dict(
         keras_name="InceptionResNetV2", scaling=(-1, 1), min_size=75, default_size=299),
     inception_v3=dict(
@@ -579,6 +593,8 @@ def _selected_model(self):
         arch = self._config["enc_architecture"]
         model = _MODEL_MAPPING.get(arch)
         model["kwargs"] = self._model_kwargs.get(arch, {})
+        if arch.startswith("efficientnet_v2"):
+            model["kwargs"]["include_preprocessing"] = False
         return model
 
     def __call__(self):

diff --git a/plugins/train/model/phaze_a_defaults.py b/plugins/train/model/phaze_a_defaults.py
@@ -52,8 +52,10 @@
 if get_backend() != "amd":
     _ENCODERS.extend(["efficientnet_b0", "efficientnet_b1", "efficientnet_b2", "efficientnet_b3",
                       "efficientnet_b4", "efficientnet_b5", "efficientnet_b6", "efficientnet_b7",
-                      "mobilenet_v3_large", "mobilenet_v3_small", "resnet50_v2", "resnet101",
-                      "resnet101_v2", "resnet152", "resnet152_v2"])
+                      "efficientnet_v2_b0", "efficientnet_v2_b1", "efficientnet_v2_b2",
+                      "efficientnet_v2_b3", "efficientnet_v2_l", "efficientnet_v2_m",
+                      "efficientnet_v2_s", "mobilenet_v3_large", "mobilenet_v3_small",
+                      "resnet50_v2", "resnet101", "resnet101_v2", "resnet152", "resnet152_v2"])
 _ENCODERS = sorted(_ENCODERS)
 
 
@@ -143,6 +145,13 @@
              "each variant is: b0: 224px, b1: 240px, b2: 260px, b3: 300px, b4: 380px, b5: 456px, "
              "b6: 528px, b7 600px. Ref: Rethinking Model Scaling for Convolutional Neural "
              "Networks (2020): https://arxiv.org/abs/1905.11946"
+             "\n\tefficientnet_v2: [Tensorflow 2.8+ only] EfficientNetV2 is the follow up to "
+             "efficientnet. It has numerous variants (B0 - B3 and Small, Medium and Large) that "
+             "increases the model width, depth and dimensional space at each step. The minimum "
+             "input resolution is 32px for all variants. The maximum input resolution for each "
+             "variant is: b0: 224px, b1: 240px, b2: 260px, b3: 300px, s: 384px, m: 480px, l: "
+             "480px. Ref: EfficientNetV2: Smaller Models and Faster Training (2021): "
+             "https://arxiv.org/abs/2104.00298"
              "\n\tfs_original: (32px - 160px). A configurable variant of the original facewap "
              "encoder. ImageNet weights cannot be loaded for this model. Additional parameters "
              "can be configured with the 'fs_enc' options. A version of this encoder is used in "