Trained an diffusion model that runs on 264KB of RAM [P]

Hardware

I recently bought a Shrike lite, which has got 264KB of SRAM. I decided to train an image generation model that generates 32*32 pixel images.

The microcontroller also has an FPGA onboard. I used it to create two parallel INT8 MAC engines with 16-bit accumulation to speed up calculations. However, the system soon hit a memory wall due to the high number of I/O operations.

Performance

This meant that the system with parallel MAC engines ran slower than the MCU-only model (~220 seconds per image vs ~70 seconds per image). It was still a fun project that I enjoyed messing around with.

A lot of the images looked weird and noisy because of the heavy quantization and memory limits, but some of them came out cool.

Full case study here. edit: added link that leads straight to the case study.