Code for the paper "BERT Loses Patience: Fast and Robust Inference with Early Exit".
NEWS: We now have a better and tidier implementation integrated into Hugging Face transformers!
If you use this code in your research, please cite our paper:
@inproceedings{zhou2020bert,
author = {Zhou, Wangchunshu and Xu, Canwen and Ge, Tao and McAuley, Julian and Xu, Ke and Wei, Furu},
booktitle = {Advances in Neural Information Processing Systems},
pages = {18330--18341},
publisher = {Curran Associates, Inc.},
title = {BERT Loses Patience: Fast and Robust Inference with Early Exit},
url = {https://proceedings.neurips.cc/paper/2020/file/d4dd111a4fd973394238aca5c05bebe3-Paper.pdf},
volume = {33},
year = {2020}
}Our code is built on huggingface/transformers. To use our code, you must clone and install huggingface/transformers.
You can fine-tune a pretrained language model and train the internal classifiers by configuring and running finetune_bert.sh and finetune_albert.sh .
You can inference with different patience settings by configuring and running patience_infer_albert.sh and patience_infer_bert.sh.
If you'd like to contribute and add more tasks (only GLUE is available at this moment), please submit a pull request and contact me. Also, if you find any problem or bug, please report with an issue. Thanks!