Skip to content

U-Net ESP32-S3 deployability: 2D-tiled decoder fits internal SRAM where TFLM OOMs - #16

Merged
asteinh merged 2 commits into
developfrom
feature/unet-esp32s3
Aug 17, 2026
Merged

U-Net ESP32-S3 deployability: 2D-tiled decoder fits internal SRAM where TFLM OOMs#16
asteinh merged 2 commits into
developfrom
feature/unet-esp32s3

Conversation

@asteinh

@asteinh asteinh commented Aug 16, 2026

Copy link
Copy Markdown
Member

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

  • QDQ reconstruction (cortex-m-deployability/tools/tflite_to_qdq_onnx.py): TransposeConv,
    Concatenation, and ResizeNearestNeighbor handlers, so a decoder .tflite reconstructs to a
    matched 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 resolver (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 the
    matched ONNX; the int8 output is non-degenerate.
  • Compile and host parity: -m 232K+6M fits (scheduled peak 232 KiB, 9.5% of the 2.38 MiB naive
    peak), with 4 of 4 ConvTranspose decoder stages carrying 2D tile plans. scripts/host_parity_unet.py
    asserts 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).
  • Bench matrix: a tigris_unet_i8_espnn cell (int8_atol=4) and a tflm_unet_i8 cell
    (expected_status ARENA_TOO_SMALL); both ESP-IDF firmwares build.
  • Housekeeping: re-pinned core-versions.json to the current develop commits; added a
    compile_tigris_plan helper so the prep scripts import cleanly; fixed a rank-4 reference-layout
    conversion (ONNX Runtime emits NCHW, the runtime emits NHWC) for spatial outputs.

Testing

Host unit tests and the tflm-esp32s3 test suite pass. Both firmwares build for esp32s3. The
U-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 ConvTranspose
    decoder stages) plus co-tiled skip chains, sram_actual_kb=232 (fits internal
    SRAM), latency_mean_ms=13733.94. That latency is the s8_ref reference-kernel
    path (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; peak
    activation 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.

@asteinh
asteinh force-pushed the feature/unet-esp32s3 branch from 32697a4 to 63569f3 Compare August 17, 2026 05:48
@asteinh
asteinh merged commit ba05bdd into develop Aug 17, 2026
2 checks passed
@asteinh
asteinh deleted the feature/unet-esp32s3 branch August 17, 2026 13:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant