diff --git a/.dockerignore b/.dockerignore index e88a421..c4c4f5d 100644 --- a/.dockerignore +++ b/.dockerignore @@ -6,3 +6,8 @@ train-input.json tensorboard_logs/ test/ __pycache__/ +scripts/ +.ei-block-config +.ei-project-config.json +out-*/ +out*.log diff --git a/.gitignore b/.gitignore index 13e1ff8..9dd8767 100644 --- a/.gitignore +++ b/.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/ diff --git a/Dockerfile b/Dockerfile index 8121e18..e4245d3 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,32 +1,5 @@ -# 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 \ @@ -34,21 +7,14 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \ 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 . ./ diff --git a/README.md b/README.md index 1ca2eaa..28c95da 100644 --- a/README.md +++ b/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)** ``` diff --git a/parameters.json b/parameters.json index 1141bdc..6b0bc86 100644 --- a/parameters.json +++ b/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" diff --git a/train.py b/train.py index 3f6ad06..2bcc9fc 100644 --- a/train.py +++ b/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))