JaxPruner: A concise library for sparsity research
Abstract
This paper introduces JaxPruner, an open-source JAX-based pruning and sparse training library for machine learning research. JaxPruner aims to accelerate research on sparse neural networks by providing concise implementations of popular pruning and sparse training algorithms with minimal memory and latency overhead. Algorithms implemented in JaxPruner use a common API and work seamlessly with the popular optimization library Optax, which, in turn, enables easy integration with existing JAX based libraries. We demonstrate this ease of integration by providing examples in four different codebases: Scenic, t5x, Dopamine and FedJAX and provide baseline experiments on popular benchmarks.
Keywords
Cite
@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}
}
Comments
Jaxpruner is hosted at http://github.com/google-research/jaxpruner