🚀 Splitting Bottlenecks: Memory-Aware Neural Architecture Search for Multi-Branch TinyML

---

📝 Overview

This project develops DupNAS, a framework that integrates neural architecture search and multi-branch tensor splitting to find high-accuracy networks that are deployable on severely memory-constrained tiny MCU-class devices. To enable this integration, DupNAS shrinks the vast splitting configuration space of multi-branch networks into a smaller set of memory-optimized configurations and explores them in a lightweight manner to resolve memory bottlenecks.

DupNAS is implemented in PyTorch on top of the popular TinyNAS framework for MCUs, with extensions to support multi-branch networks and integration with our splitting strategy. The resulting network solutions are deployed on an STM32F746 microcontroller (ARM Cortex-M7 CPU, 216 MHz, 320 KB VM, and 1 MB NVM) running the TensorFlow Lite Micro (TFLite-Micro) inference engine. We evaluate DupNAS on three vision-based TinyML network families, namely MobileNetV2, ShuffleNetV2, and InceptionV3, under different memory constraints. All networks are trained on the ImageNet-100 dataset. DupNAS is compared against two existing splitting strategies, TinyTS and PatchTS.

Overview of the DupNAS framework

📌 File Structure

Below is a brief description of the main directories and files in this repository.

• /DupNAS/NASBase/duplication implements the DupNAS splitting configuration exploration algorithm, which is invoked during the NAS process.
• /DupNAS/NASBase/ss_optimization contains the NAS architecture space optimization component, adapted from TinyNAS.
• /DupNAS/NASBase/evo_search contains the NAS evolutionary search component, adapted from TinyNAS.
• /DupNAS/NASBase/model defines the architecture search space, including supernet and subnet architecture definitions.
• /DupNAS/settings provides the evaluation settings for different datasets and baselines.
• /DupNAS/settings.py defines the global NAS settings.
• /Inference/Model-converter/ provides the model converter, which applies the splitting configuration to a network solution to generate a split network that can be deployed on TFLite-Micro.
• /Inference/Tflm-engine/ provides the build process for TFLite-Micro.
• /assets/DupNAS_paper_data.xlsx contains the experimental results presented in the paper.
• /assets/models/ contains the models used for evaluation and deployment.

---

🧭 Getting Started

💡 Prerequisites

• Python 3.9.
• Install the required Python packages listed in requirements.txt with: python3.9 -m pip install -r requirements.txt
• Anaconda (recommended for managing Python environments).
• ImageNet-100 dataset.
• STM32F746G-DISCO deployment device.
• STM32H747I-DISCO deployment device.
• STM32CubeIDE 2.1.1 development tools for the STM32.
• STM32CubeMX 6.17.0 development tools for the STM32.
• STM32Cube AI Studio 1.2.0 development tools and inference engine for the STM32.
• TensorFlow Lite Micro inference engine.

🔧 Setup and running DupNAS

1. Download or clone this repository.
2. Create and activate a Python environment, then install the required dependencies.

   python3.9 -m pip install -r requirements.txt

3. Configure access to the ImageNet-100 dataset
   • Open the ImageNet-100 Kaggle page.
   • Generate a Kaggle API token.
   • Export the token:

     export KAGGLE_API_TOKEN="<your-token>"

   • Store the token in the Kaggle configuration directory:

     mkdir -p ~/.kaggle
     echo "<your-token>" > ~/.kaggle/access_token
     chmod 600 ~/.kaggle/access_token

   The dataset is downloaded and prepared automatically when DupNAS is run for the first time stage 2.
4. Copy the configuration file for the target architecture and run DupNAS:

   cd DupNAS/
   cp settings/settings-<arc>.py settings.py
   
   python3.9 -m NASBase.run_nas \
     --stages <stage> \
     --arc <arc> \
     --dataset IMAGE100 \
     --mode <mode> \
     --vmsize <vmsize> \
     --suffix <suffix> \
     --no-rlogger

📝 Arguments

Options	Description	Candidate Values
--stages	Number of NAS stages: ssopt, training, evosearch, fine-tuning	1, 2, 3, 4
--arc	network architecture family	mbv2, shuffle, incept
--dataset	Dataset used for search	IMAGE100
--mode	splitting strategy	dupnas, tinyts, patchts, nots
--vmsize	VM constraint in KB	96, 128, 256
--suffix	Experiment suffix for naming outputs. Use the same suffix for all four stages.	user-defined string

5. After completing Stages 1–4, extract the selected network configuration, the final solution is saved in /DupNAS/NASBase/train_log/<suffix>_best_solution.json
6. ONNX generation and tensor-splitting conversion are integrated into the model-converter workflow. Continue with steps in the next section.

✂️ Setup and running the model converter

1. Copy the configurations "supernet config" and "subnet_choice_per_blk" from: <suffix>_best_solution.json, into: /DupNAS/NASBase/spec_model_<arc>.txt.
2. Go to /DupNAS/, then run the corresponding script to generate the selected ONNX models:

   bash gen_selected_shuffle.sh
   bash gen_selected_mbv2.sh
   bash gen_selected_incept.sh

3. The generated models are saved in the following locations:
   • Original selected ONNX models: /DupNAS/genonnx/<arc>/
   • Tensor-split ONNX models: /Inference/Model-converter/ts_converted/<arc>/
4. Go to the /Inference/onnx-to-tflite/ directory and run the corresponding conversion command:

   bash convert.sh shuffle
   bash convert.sh mbv2
   bash convert.sh incept

   The conversion uses onnx2tf to generate integer-quantized TensorFlow Lite models, such as:

   xxx_full_integer_quant.tflite
   

   The generated TFLite models are saved under /Inference/onnx-to-tflite/outputs/

⚙️ Setup and building TFLite-Micro

1. Copy the converted TFLite model (xxx_full_integer_quant.tflite) into Tflm-engine/src/models.
2. Follow Tflm-engine/README.md to build the TFLite-Micro static library (libtensorflow-microlite.a).
3. Add the generated static library to your STM32CubeIDE project settings. Then include Tflm-engine/src/tflm_main.h and call tflm_main_xxx to run inference for the target model.

For more information, please refer to Tflm-engine/README.md.

---

🧩 Evaluation

Search Space Configuration

Level	Option	MobileNetV2	ShuffleNetV2	InceptionV3
Backbone	Branches	1	2	4
Backbone	Blocks	4	3	3
Supernet	Input resolution	32, 64, 96, 128	32, 64, 96, 128	32, 64, 96, 128
Supernet	Width multiplier	0.25, 0.5, 0.75, 1.0	0.2, 0.5, 0.8	0.25, 0.5, 0.75, 1.0
Block	Kernel size	3, 5	1, 3, 5, 7	3, 5, 7
Block	Expansion / stride	expansion: 3, 4, 6	stride: 1, 2	stride: 1, 2
Block	Layers	1, 2, 3	1, 2, 3	1, 2, 3

Accuracy

Below are the networks found by DupNAS, TinyTS, and PatchTS.

Model	VM	DupNAS	TinyTS	PatchTS
MobileNetV2	96 KB	58.40%	52.88%	51.36%
	128 KB	62.08%	56.88%	56.48%
	256 KB	62.64%	61.76%	62.64%
ShuffleNetV2	96 KB	61.36%	56.8%	54.24%
	128 KB	62.96%	59.76%	58.72%
	256 KB	65.76%	64.96%	60.48%
InceptionV3	96 KB	61.84%	45.68%	54.00%
	128 KB	64.16%	57.84%	58.64%
	256 KB	68.24%	64.88%	67.36%