U-Net ESP32-S3 deployability: 2D-tiled decoder fits internal SRAM where TFLM OOMs - #16
Merged
Conversation
asteinh
force-pushed
the
feature/unet-esp32s3
branch
from
August 17, 2026 05:48
32697a4 to
63569f3
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a U-Net segmentation model to the ESP32-S3 deployability suite. A 256x256 int8 U-Net
(Conv2DTranspose-upsampled decoder with concat skips) compiles with TiGrIS and 2D-tiles its
high-resolution decoder into internal SRAM while spilling the skip tensors to PSRAM. A TFLite
Micro build of the same model with a 256 KB internal-SRAM arena fails AllocateTensors, because
the peak activation is 1.19 MiB.
What changed
cortex-m-deployability/tools/tflite_to_qdq_onnx.py): TransposeConv,Concatenation, and ResizeNearestNeighbor handlers, so a decoder
.tflitereconstructs to amatched QDQ ONNX. RESIZE_BILINEAR is rejected (TFLite integer bilinear is not bit-exactly
reproducible). A unit test reconstructs a three-op model and matches the TFLite int8 interpreter
within 1 LSB.
tflm-esp/main/main.cc): registered TransposeConv/Concatenation/ResizeNearestNeighbor so a decoder loads and fails on arena size rather than a missing op.
models/prepare.py build_keras_unet: a 256x256x3 -> 256x256x8 int8 U-Net (encoder 4x stride-2,Conv2DTranspose decoder + concat skips). Produces the
.tflite, the TFLM C header, and thematched ONNX; the int8 output is non-degenerate.
-m 232K+6Mfits (scheduled peak 232 KiB, 9.5% of the 2.38 MiB naivepeak), with 4 of 4 ConvTranspose decoder stages carrying 2D tile plans.
scripts/host_parity_unet.pyasserts the tiled runtime output equals the untiled output bit-exactly and tracks ONNX Runtime
within the deep-int8 tolerance (4 LSB, within the measured ORT-vs-TFLite oracle spread and
consistent with the existing DS-CNN cell).
tigris_unet_i8_espnncell (int8_atol=4) and atflm_unet_i8cell(expected_status ARENA_TOO_SMALL); both ESP-IDF firmwares build.
core-versions.jsonto the current develop commits; added acompile_tigris_planhelper so the prep scripts import cleanly; fixed a rank-4 reference-layoutconversion (ONNX Runtime emits NCHW, the runtime emits NHWC) for spatial outputs.
Testing
Host unit tests and the
tflm-esp32s3test suite pass. Both firmwares build for esp32s3. TheU-Net benchmark cells run through the standard bench flow.
On-device (real ESP32-S3)
Captured on a physical ESP32-S3 (512 KB internal SRAM, 8 MB PSRAM):
tigris_unet_i8_espnn: runs.stages_tiled=4(the 2D-tiled ConvTransposedecoder stages) plus co-tiled skip chains,
sram_actual_kb=232(fits internalSRAM),
latency_mean_ms=13733.94. That latency is the s8_ref reference-kernelpath (ConvTranspose and Concat are not esp-nn-accelerated), not an accelerated
headline. Device output matches the reference within the int8 tolerance
(max abs diff 3, atol 4).
tflm_unet_i8:status=ARENA_TOO_SMALL(256 KB internal-SRAM arena; peakactivation 1.19 MiB).
So the same U-Net runs on the ESP32-S3 by tiling its decoder into internal SRAM
and spilling skips to PSRAM, where a TFLM internal-SRAM arena cannot allocate.