This repo is tested on Python 3.7.7, PyTorch 1.4.0, and Cuda 10.1. Using a virtulaenv or conda environemnt is recommended, for example:

conda install pytorch==1.4.0 torchvision cudatoolkit=10.1 -c pytorch

After installing the required environment and cloning this repo, install the following requirements:

pip install -r ./requirements.txt
pip install -r ./examples/requirements.txt

Download GLUE dataset by

python download_glue_data.py --data_dir data --tasks all

Scripts are in the scripts folder, which corresponds to RomeBERT - SD only.

This is for fine-tuning RomeBERT - SD only models.

This is for evaluating each exit layer for fine-tuned RomeBERT - SD only models.

This is for evaluating fine-tuned RomeBERT - SD only models, given a number of different early exit entropy thresholds.

This is for creating output tsv file for test split for fine-tuned RomeBERT - SD only models.

Scripts are in the new_scripts folder, which corresponds to RomeBERT - SD+GR.

This is for fine-tuning RomeBERT - SD+GR models.

This is for evaluating each exit layer for fine-tuned RomeBERT - SD+GR models.

This is for evaluating fine-tuned RomeBERT - SD+GR models, given a number of different early exit entropy thresholds.

This is for creating output tsv file for test split for fine-tuned RomeBERT - SD+GR models.

@misc{geng2021romebert,
    title={RomeBERT: Robust Training of Multi-Exit BERT},
    author={Shijie Geng and Peng Gao and Zuohui Fu and Yongfeng Zhang},
    year={2021},
    eprint={2101.09755},
    archivePrefix={arXiv},
    primaryClass={cs.CL}
}