PyTorch custom learning block (ExecuTorch)

A minimal Edge Impulse custom learning block that trains a small PyTorch CNN on image data and exports:

• model.onnx — consumed by Edge Impulse (auto-converted to TFLite for deployment).
• model.pte — an optional ExecuTorch program for native PyTorch on-device inference.

This is the MVP starter sample for the ExecuTorch end-to-end series.

Files

File	Purpose
parameters.json	Block metadata and the parameters shown in Studio (epochs, learning rate, batch size, .pte export flag).
model.py	SmallCNN, a compact resolution-independent classifier.
train.py	Training entrypoint. Loads the NumPy splits, trains, and exports ONNX (+ .pte).
Dockerfile	Builds the training container.
requirements.txt	Pinned Python dependencies.
tests/test_model.py	Smoke tests for model shapes and data handling.

How it works

Edge Impulse provides image data already pre-processed and pixel-scaled, in NHWC float32 arrays. The block:

1. Loads X_split_{train,test}.npy and Y_split_{train,test}.npy from --data-directory.
2. Transposes NHWC → NCHW for PyTorch (Edge Impulse handles the reverse on-device because we output ONNX).
3. Trains SmallCNN with Adam + cross-entropy.
4. Writes model.onnx (and model.pte when --export-pte is set) to --out-directory.

Test locally

Download data from a project that has an image impulse, then run the block:

# Pull the latest processed data into ./input
edge-impulse-blocks runner --download-data input/

# Build and run the training container
docker build -t pytorch-custom-deploy .
docker run --rm -v "$PWD":/app pytorch-custom-deploy \
    --data-directory /app/input \
    --out-directory /app/out \
    --epochs 30 --learning-rate 0.001 --batch-size 32 --export-pte

After training you'll find out/model.onnx (and out/model.pte).

Run the unit tests without Docker:

pip install -r requirements.txt pytest
python -m pytest

Push to Edge Impulse

edge-impulse-blocks init   # first time only
edge-impulse-blocks push

The block then appears under Create impulse → Add learning block in Studio.

Walkthrough in Studio

The screenshots below are from a 3-class image project (lamp / plant / unknown) that uses this block end to end.

Step	Screenshot
Create impulse (Image processing + this learning block)
Training (93.3% accuracy on the validation set)
Model testing
Deployment (incl. ExecuTorch .pte export target)

Notes

• imageInputScaling is set to 0..1 in parameters.json. Match this to how you want pixels scaled; Edge Impulse applies it before your block sees the data.
• The .pte file is not used by the Edge Impulse SDK. It's included so you can carry the same trained weights into a native ExecuTorch runtime (see the executorch-deploy block).
• Pin torch and executorch to compatible versions. If ExecuTorch isn't installed, .pte export is skipped with a warning instead of failing the job.