JaxPruner: A concise library for sparsity research
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.
Code (1)
Methods 이 논문이 사용한 방법론
Similar Papers 제목 키워드 기반
Pixyz: a Python library for developing deep generative models
With the recent rapid progress in the study of deep generative models (DGMs), there is a need for a framework that can implement them in a simple and generic way. In this research, we focus on two features of DGMs: (1) d…
Probabilistic ProgrammingAutomated discovery of interpretable hyperelastic material models for human brain tissue with EUCLID
We propose an automated computational algorithm for simultaneous model selection and parameter identification for the hyperelastic mechanical characterization of human brain tissue. Following the motive of the recently p…
Model SelectionregressionGPflow: A Gaussian process library using TensorFlow
GPflow is a Gaussian process library that uses TensorFlow for its core computations and Python for its front end. The distinguishing features of GPflow are that it uses variational inference as the primary approximation …
Gaussian ProcessesGPUsoftware testingVariational Inferencefastai: A Layered API for Deep Learning
fastai is a deep learning library which provides practitioners with high-level components that can quickly and easily provide state-of-the-art results in standard deep learning domains, and provides researchers with low-…
Deep LearningGPUFuzzyLogic.jl: a Flexible Library for Efficient and Productive Fuzzy Inference
This paper introduces \textsc{FuzzyLogic.jl}, a Julia library to perform fuzzy inference. The library is fully open-source and released under a permissive license. The core design principles of the library are: user-frie…