JaxPruner：面向稀疏性研究的简洁库
机器学习
2023-12-20 v3 软件工程
摘要
本文介绍了JaxPruner，一个基于JAX的开源剪枝与稀疏训练库，用于机器学习研究。JaxPruner旨在通过以最小的内存和延迟开销提供主流剪枝与稀疏训练算法的简洁实现，来加速稀疏神经网络的研究。JaxPruner中实现的算法使用通用API，并与流行的优化库Optax无缝协作，这反过来使其能轻松集成到现有的基于JAX的库中。我们通过在四个不同的代码库：Scenic、t5x、Dopamine和FedJAX中提供示例来展示这种易于集成的特性，并在流行的基准测试上提供基线实验。
引用
@article{arxiv.2304.14082,
title = {JaxPruner: A concise library for sparsity research},
author = {Joo Hyung Lee and Wonpyo Park and Nicole Mitchell and Jonathan Pilault and Johan Obando-Ceron and Han-Byul Kim and Namhoon Lee and Elias Frantar and Yun Long and Amir Yazdanbakhsh and Shivani Agrawal and Suvinay Subramanian and Xin Wang and Sheng-Chun Kao and Xingyao Zhang and Trevor Gale and Aart Bik and Woohyun Han and Milen Ferev and Zhonglin Han and Hong-Seok Kim and Yann Dauphin and Gintare Karolina Dziugaite and Pablo Samuel Castro and Utku Evci},
journal= {arXiv preprint arXiv:2304.14082},
year = {2023}
}
备注
Jaxpruner is hosted at http://github.com/google-research/jaxpruner