distributed_transformer_inference

This repos is the accompanying repo of our paper Going beyond the Edge: Distributed Inference of Transformer Models on Ultra-low-power Wireless Devices

Installation

The code has been tested with Python 3.12. Install all packages with

pip install -r requirements.txt

Entrance points

train.py

Starts the training. We use hydra (https://hydra.cc/) for configuring our training runs. The parameters can be found in parameters/. The folder parameters/run contains the dataset-specific parameter changes.

Example usage:

python train.py "+run=cat_vs_dog" "dataset.dataset_base_path=<path to dataset>" "root_dir=<path to logging files>"

export_to_hardware.py

Exports a given neural network to hardware code.

Example usage:

python export_to_hardware.py "+run=cat_vs_dog" "dataset.dataset_base_path=<path to dataset>" "root_dir=<path to logging files>"

hardware_implementation/

Our code for the hardware implementation. Open hardware_implementation/firmware/firmware.emProject in Segger embedded studio (https://www.segger.com/products/development-tools/embedded-studio/). You can compile and flash the code using hardware_implementation/build_and_flash.bash.

The implementation of the layers can be found in hardware_implementation/firmware/distributed_inference/.