ProTran is a tool which can be used to generate and evaluate different Transformer architectures on a diverse set of embedded platforms for various natural language processing tasks. This repository uses the FlexiBERT framework (jha-lab/txf_design-space) to obtain the design space of flexible and heterogeneous Transformer models.
Supported platforms:
- Linux on x86 CPUs with CUDA GPUs (tested on AMD EPYC Rome CPU, Intel Core i7-8650U CPU and Nvidia A100 GPU).
- Apple M1 and M1-Pro SoC on iPad and MacBook Pro respectively.
- Broadcom BCM2711 SoC on Raspberry Pi 4 Model-B.
- Intel Neural Compute Stick v2.
- Nvidia Tegra X1 SoC on Nvidia Jetson Nano 2GB.
git clone --recurse-submodules https://github.com/shikhartuli/protran.git
cd protran
For Unix-like platforms, download and install conda through miniforge:
curl -L -O "https://github.com/conda-forge/miniforge/releases/latest/download/Mambaforge-$(uname)-$(uname -m).sh"
bash Mambaforge-$(uname)-$(uname -m).sh
The python environment setup is based on conda. The script below creates a new environment named txf_design-space or updates an existing environment:
source env_step.sh
To test the installation, you can run:
cd txf_design-space
python check_install.py
cd ..
All training scripts use bash and have been implemented using SLURM. This will have to be setup before running the experiments.
Shikhar Tuli. For any questions, comments or suggestions, please reach me at stuli@princeton.edu.
Cite our work using the following bitex entry:
@article{tuli2023edgetran,
title={{EdgeTran}: Device-Aware Co-Search of Transformers for Efficient Inference on Mobile Edge Platforms},
author={Tuli, Shikhar and Jha, Niraj K},
journal={IEEE Transactions on Mobile Computing},
year={2023}
}BSD-3-Clause. Copyright (c) 2023, Shikhar Tuli and Jha Lab. All rights reserved.
See License file for more details.