Adding custom learning blocks

Enterprise customers can add fully custom learning models in Edge Impulse. These models can be represented in any deep learning framework as long as it can output TFLite files.

Only available for enterprise customers

Organizational features are only available for enterprise customers. View our pricing for more information.

Custom learning blocks for organizations go beyond project's expert mode, which lets you to design your own neural network architectures on a project by project basis, but has some caveats:

• You can't load pretrained weights from expert mode, and

• Your expert mode model needs to be representable in Keras / TensorFlow.

	Custom Learning Blocks (Organizations)	Expert Mode (Projects)
Load pretrained weights	✅	❌
Use any ML framework to define your model	✅	Keras only

This tutorial describes how to build these models. Alternatively, we've put together two example projects, which bring these models into Edge Impulse:

• YOLOv5 - brings a YOLOv5 transfer learning model (trained with PyTorch) into Edge Impulse

• Keras - shows how to bring custom Keras blocks into Edge Impulse.

• PyTorch - shows how to bring custom PyTorch blocks into Edge Impulse.

Prerequisites

• The Edge Impulse CLI

  • If you receive any warnings that's fine. Run edge-impulse-blocks afterwards to verify that the CLI was installed correctly.

• Docker desktop - Custom learning models use Docker containers, a virtualization technique which lets developers package up an application with all dependencies in a single package.

Data formats

To bring custom learning models into Edge Impulse you'll need to encapsulate your training pipeline into a container. This container takes in the training data, trains the model and then spits out TFLite files.

To see the data that will be passed into the container:

1. Create a project in Edge Impulse and add your data.

2. Under Create impulse add the DSP block that you'll use (e.g. 'Image' for images, or 'Spectrogram' for audio) and a random neural network block.

3. Generate features for the DSP block.

4. Now go to Dashboard, and under 'Download block data' grab the two items marked 'Training data' / 'Training labels'.

Downloading DSP block output

This data is in the following formats:

• Training data: a numpy matrix with one sample per row containing the output of the DSP block (use np.load('X_train_features.npy') to see).

• Training labels

  • If you're using object detection: a JSON file (despite the '.npy' extension) containing structured information about the bounding boxes. Every item in the samples array maps to one row in the data matrix. The label is the index of the class, and is 1-based.

  • If you're not using object detection: a numpy matrix with one sample per row, and the first column of every row is the index of the class. The last three columns are the sampleId and the start / end time of the sample (when going through time series). During training you can typically discard this.

This data is passed into the container as files, located here:

• Data: /home/X_train_features.npy

• Labels: /home/y_train.npy

After training you need to output a TFLite file. You need to write these here:

• Float32 (unquantized) model: /home/model.tflite

• Int8 quantized model with int8 inputs: /home/model_quantized_int8_io.tflite

Wrapping your training scripts in a Docker container

Docker containers are a virtualization technique which lets developers package up an application with all dependencies in a single package. To train your custom model you'll need to wrap all the required packages, your scripts, and (if you use transfer learning) your pretrained weights into this container. When running in Edge Impulse the container does not have network access, so make sure you don't download dependencies while running (fine when building the container).

A typical Dockerfile might look like:

# syntax = docker/dockerfile:experimental
FROM ubuntu:20.04
WORKDIR /app

ARG DEBIAN_FRONTEND=noninteractive

# Install base packages (like Python and pip)
RUN apt update && apt install -y curl zip git lsb-release software-properties-common apt-transport-https vim wget python3 python3-pip
RUN python3 -m pip install --upgrade pip==20.3.4

# Copy Python requirements in and install them
COPY requirements.txt ./
RUN pip3 install -r requirements.txt

# Copy the rest of your training scripts in
COPY . ./

# And tell us where to run the pipeline
ENTRYPOINT ["python3", "-u", "train.py"]

It's important to create an ENTRYPOINT at the end of the Dockerfile to specify which file to run.

Arguments during training

The train script will receive the following arguments:

• --epochs <epochs> - number of epochs to train (e.g. 50).

• --learning-rate <lr> - learning rate (e.g. 0.001).

• --validation-set-size <size> - size of the validation set (e.g. 0.2 for 20% of total training set).

• --input-shape <shape> - shape of the training data (e.g. (320,320,3) for a 320x320 RGB image).

Running the training script

To run the container, first create a home folder in your script directory and copy the training data / training labels here (named X_train_features.npy and y_train.npy). Then build and run the container:

$ docker build -t mycontainer .
$ docker run --rm -v $PWD/home:/home mycontainer --epochs 50 --learning-rate 0.001 --validation-set-size 0.2 --input-shape "(320,320,3)"

This should train the model and spit out .tflite files.

Pushing the block to Edge Impulse

If your block works you can bring it into Edge Impulse via:

$ edge-impulse-blocks init
$ edge-impulse-blocks push

To edit the block, go to your organization, Custom blocks > Transfer learning models.

The block is now available for every Edge Impulse project under your organization.

Selecting the new learning model.

Object detection output layers

Unfortunately object detection models typically don't have a standard way to go from neural network output layer to bounding boxes. Currently we support the following types of output layers:

• MobileNet SSD

• Edge Impulse FOMO

• YOLOv5

If you have an object detection model with a different output layer then please contact your user success engineer with an example on how to interpret the output, and we can add it.