Add support for EfficientNetV2

.dockerignore

@@ -6,3 +6,8 @@ train-input.json
 tensorboard_logs/
 test/
 __pycache__/
+scripts/
+.ei-block-config
+.ei-project-config.json
+out-*/
+out*.log

.gitignore

@@ -10,3 +10,8 @@ transfer-learning-weights/*.h5
 train_input.json
 tensorboard_logs/
 test/
+.ei-project-config.json
+scripts/node_modules/
+out-*/
+out*.log
+.vscode/

Dockerfile

@@ -1,54 +1,20 @@
-# syntax = docker/dockerfile:experimental@sha256:3c244c0c6fc9d6aa3ddb73af4264b3a23597523ac553294218c13735a2c6cf79
-ARG UBUNTU_VERSION=24.04
-
-ARG ARCH=
-ARG CUDA=12.9.1
-ARG CUDA_SHORT=12.9
-ARG CUDA_PACKAGE_VERSION=12-9
-ARG CUDA_FLAVOR=base
-FROM nvidia/cuda${ARCH:+-$ARCH}:${CUDA}-${CUDA_FLAVOR}-ubuntu${UBUNTU_VERSION} as base
-ARG CUDA
-ARG CUDA_SHORT
-ARG CUDA_PACKAGE_VERSION
-ENV DEBIAN_FRONTEND=noninteractive
-
-WORKDIR /app
-
-# Install Python, pip, and dos2unix (as when you check out install_cuda.sh on Windows it converts to CRLF which bash does not like in the next step)
-RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
-    apt-get update && apt-get install -y --no-install-recommends \
-        python3 python3-pip dos2unix && \
-    rm -rf /var/lib/apt/lists/*
-
-# Install NVIDIA CUDA/cuDNN runtime libraries needed by TensorFlow on x86.
-COPY dependencies/install_cuda.sh ./install_cuda.sh
-RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
-    dos2unix ./install_cuda.sh && \
-    /bin/bash ./install_cuda.sh && \
-    rm install_cuda.sh && \
-    rm -rf /var/lib/apt/lists/*
+# Simple Ubuntu 24.04 base image with Python3.12 and CUDA setup already (for GPU training)
+FROM public.ecr.aws/z9b3d4t5/ei-custom-ml-block-base:v1.95.5-test-9e8dfa82
 
 # Add other system dependencies
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     apt-get update && apt-get install -y --no-install-recommends \
         wget && \
     rm -rf /var/lib/apt/lists/*
 
-# Download weights, mirrored from https://github.com/Runist/image-classifier-keras/releases
-RUN mkdir -p /weights && \
-    cd /weights && \
-    wget https://cdn.edgeimpulse.com/pretrained-weights/efficientnet/efficientnetb0_notop.h5 && \
-    wget https://cdn.edgeimpulse.com/pretrained-weights/efficientnet/efficientnetb1_notop.h5 && \
-    wget https://cdn.edgeimpulse.com/pretrained-weights/efficientnet/efficientnetb2_notop.h5 && \
-    wget https://cdn.edgeimpulse.com/pretrained-weights/efficientnet/efficientnetb3_notop.h5 && \
-    wget https://cdn.edgeimpulse.com/pretrained-weights/efficientnet/efficientnetb4_notop.h5 && \
-    wget https://cdn.edgeimpulse.com/pretrained-weights/efficientnet/efficientnetb5_notop.h5
-
 # Copy Python requirements in and install them (--break-system-packages is required if we don't use a venv)
 COPY requirements.txt ./
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip3 install --break-system-packages -r requirements.txt
 
+# Pre-cache ImageNet weights so transfer learning also works with --network=none.
+RUN python3 -c "import tensorflow as tf; [builder(include_top=False, pooling='avg', weights='imagenet') for builder in (tf.keras.applications.EfficientNetB0, tf.keras.applications.EfficientNetB1, tf.keras.applications.EfficientNetB2, tf.keras.applications.EfficientNetB3, tf.keras.applications.EfficientNetB4, tf.keras.applications.EfficientNetB5, tf.keras.applications.EfficientNetV2B0, tf.keras.applications.EfficientNetV2B1, tf.keras.applications.EfficientNetV2B2, tf.keras.applications.EfficientNetV2B3, tf.keras.applications.EfficientNetV2S, tf.keras.applications.EfficientNetV2M, tf.keras.applications.EfficientNetV2L)]"
+
 # Copy the rest of your training scripts in
 COPY . ./
 

README.md

@@ -1,6 +1,6 @@
 # EfficientNet base models for Edge Impulse
 
-This repository contains the code to bring EfficientNet models into Edge Impulse.
+This repository contains the code to bring EfficientNet and EfficientNetV2 models into Edge Impulse.
 
 ## Using this model
 
@@ -53,6 +53,8 @@ You run this pipeline via Docker. This encapsulates all dependencies and package
         --out-directory out/
     ```
 
+    Use the original EfficientNet values (`b0` through `b5`) for backwards compatibility, or use the prefixed EfficientNetV2 values (`v2_b0`, `v2_b1`, `v2_b2`, `v2_b3`, `v2_s`, `v2_m`, `v2_l`).
+
     **Windows (Command prompt)**
 
     ```

parameters.json

@@ -3,7 +3,7 @@
     "type": "machine-learning",
     "info": {
         "name": "EfficientNet",
-        "description": "A collection of high-end image classification models, with transfer learning weights, scaling from 4 to 28M parameters. Great for complex usecases on Linux.",
+        "description": "A collection of high-end EfficientNet and EfficientNetV2 image classification models, with transfer learning weights. Great for complex usecases on Linux.",
         "operatesOn": "image",
         "imageInputScaling": "0..255",
         "indRequiresGpu": false,
@@ -30,28 +30,56 @@
             "type": "select",
             "valid": [
                 {
-                    "label": "B0 - 4M params, 16 MB",
+                    "label": "V1 B0 - 4M params, 16 MB",
                     "value": "b0"
                 },
                 {
-                    "label": "B1 - 6.5M params, 26 MB",
+                    "label": "V1 B1 - 6.5M params, 26 MB",
                     "value": "b1"
                 },
                 {
-                    "label": "B2 - 7.7M params, 30.8 MB",
+                    "label": "V1 B2 - 7.7M params, 30.8 MB",
                     "value": "b2"
                 },
                 {
-                    "label": "B3 - 10.7M params, 42.8 MB",
+                    "label": "V1 B3 - 10.7M params, 42.8 MB",
                     "value": "b3"
                 },
                 {
-                    "label": "B4 - 17.5M params, 70 MB",
+                    "label": "V1 B4 - 17.5M params, 70 MB",
                     "value": "b4"
                 },
                 {
-                    "label": "B5 - 28.3M params, 113.2 MB",
+                    "label": "V1 B5 - 28.3M params, 113.2 MB",
                     "value": "b5"
+                },
+                {
+                    "label": "V2 B0 - 5.9M params, 24 MB",
+                    "value": "v2_b0"
+                },
+                {
+                    "label": "V2 B1 - 6.9M params, 28 MB",
+                    "value": "v2_b1"
+                },
+                {
+                    "label": "V2 B2 - 8.8M params, 35 MB",
+                    "value": "v2_b2"
+                },
+                {
+                    "label": "V2 B3 - 12.9M params, 52 MB",
+                    "value": "v2_b3"
+                },
+                {
+                    "label": "V2 S - 21.6M params, 86 MB",
+                    "value": "v2_s"
+                },
+                {
+                    "label": "V2 M - 54.4M params, 218 MB",
+                    "value": "v2_m"
+                },
+                {
+                    "label": "V2 L - 119M params, 475 MB",
+                    "value": "v2_l"
                 }
             ],
             "param": "model-size"

train.py

@@ -21,10 +21,24 @@
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 
-WEIGHTS_PREFIX = os.environ.get('WEIGHTS_PREFIX', '/weights')
+MODEL_VARIANTS = {
+    'b0': tf.keras.applications.EfficientNetB0,
+    'b1': tf.keras.applications.EfficientNetB1,
+    'b2': tf.keras.applications.EfficientNetB2,
+    'b3': tf.keras.applications.EfficientNetB3,
+    'b4': tf.keras.applications.EfficientNetB4,
+    'b5': tf.keras.applications.EfficientNetB5,
+    'v2_b0': tf.keras.applications.EfficientNetV2B0,
+    'v2_b1': tf.keras.applications.EfficientNetV2B1,
+    'v2_b2': tf.keras.applications.EfficientNetV2B2,
+    'v2_b3': tf.keras.applications.EfficientNetV2B3,
+    'v2_s': tf.keras.applications.EfficientNetV2S,
+    'v2_m': tf.keras.applications.EfficientNetV2M,
+    'v2_l': tf.keras.applications.EfficientNetV2L,
+}
 
 # Load files
-parser = argparse.ArgumentParser(description='EfficientNet B0 model in Edge Impulse')
+parser = argparse.ArgumentParser(description='EfficientNet and EfficientNetV2 models in Edge Impulse')
 parser.add_argument('--info-file', type=str, required=False,
                     help='train_input.json file with info about classes and input shape')
 parser.add_argument('--data-directory', type=str, required=True)
@@ -119,46 +133,17 @@
         enable_tensorboard=True)
 
 # model architecture
-if model_size == 'b0':
-    if use_pretrained_weights:
-        weights_path = os.path.join(WEIGHTS_PREFIX, 'efficientnetb0_notop.h5')
-        base_model = tf.keras.applications.EfficientNetB0(include_top=False, pooling='avg', weights=weights_path, classes=classes)
-    else:
-        base_model = tf.keras.applications.EfficientNetB0(include_top=False, pooling='avg', weights=None, classes=classes)
-elif model_size == 'b1':
-    if use_pretrained_weights:
-        weights_path = os.path.join(WEIGHTS_PREFIX, 'efficientnetb1_notop.h5')
-        base_model = tf.keras.applications.EfficientNetB1(include_top=False, pooling='avg', weights=weights_path, classes=classes)
-    else:
-        base_model = tf.keras.applications.EfficientNetB1(include_top=False, pooling='avg', weights=None, classes=classes)
-elif model_size == 'b2':
-    if use_pretrained_weights:
-        weights_path = os.path.join(WEIGHTS_PREFIX, 'efficientnetb2_notop.h5')
-        base_model = tf.keras.applications.EfficientNetB2(include_top=False, pooling='avg', weights=weights_path, classes=classes)
-    else:
-        base_model = tf.keras.applications.EfficientNetB2(include_top=False, pooling='avg', weights=None, classes=classes)
-elif model_size == 'b3':
-    if use_pretrained_weights:
-        weights_path = os.path.join(WEIGHTS_PREFIX, 'efficientnetb3_notop.h5')
-        base_model = tf.keras.applications.EfficientNetB3(include_top=False, pooling='avg', weights=weights_path, classes=classes)
-    else:
-        base_model = tf.keras.applications.EfficientNetB3(include_top=False, pooling='avg', weights=None, classes=classes)
-elif model_size == 'b4':
-    if use_pretrained_weights:
-        weights_path = os.path.join(WEIGHTS_PREFIX, 'efficientnetb4_notop.h5')
-        base_model = tf.keras.applications.EfficientNetB4(include_top=False, pooling='avg', weights=weights_path, classes=classes)
-    else:
-        base_model = tf.keras.applications.EfficientNetB4(include_top=False, pooling='avg', weights=None, classes=classes)
-elif model_size == 'b5':
-    if use_pretrained_weights:
-        weights_path = os.path.join(WEIGHTS_PREFIX, 'efficientnetb5_notop.h5')
-        base_model = tf.keras.applications.EfficientNetB5(include_top=False, pooling='avg', weights=weights_path, classes=classes)
-    else:
-        base_model = tf.keras.applications.EfficientNetB5(include_top=False, pooling='avg', weights=None, classes=classes)
-else:
-    print(f'Expected --model-size to be b0, b1, b2, b3, b4 or b5 (was {model_size})')
+model_builder = MODEL_VARIANTS.get(model_size)
+if model_builder is None:
+    print(f'Expected --model-size to be one of {", ".join(MODEL_VARIANTS.keys())} (was {model_size})')
     exit(1)
 
+base_model = model_builder(
+    include_top=False,
+    pooling='avg',
+    weights='imagenet' if use_pretrained_weights else None,
+    classes=classes)
+
 if use_pretrained_weights:
     # What percentage of the base model's layers we will fine tune
     fine_tune_from = math.ceil(len(base_model.layers) * (freeze_percentage_of_layers / 100))