TorchOpt: An Efficient Library for Differentiable Optimization

Differentiable optimization algorithms often involve expensive computations of various meta-gradients. To address this, we design and implement TorchOpt, a new PyTorch-based differentiable optimization library. TorchOpt provides an expressive and unified programming interface that simplifies the implementation of explicit, implicit, and zero-order gradients. Moreover, TorchOpt has a distributed execution runtime capable of parallelizing diverse operations linked to differentiable optimization tasks across CPU and GPU devices. Experimental results demonstrate that TorchOpt achieves a 5.2× training time speedup in a cluster. TorchOpt is open-sourced at https://github.com/metaopt/torchopt and has become a PyTorch Ecosystem project. [abs] [ pdf ][ bib ] [ code ] &copy JMLR 2023. ( edit, beta )

Authors

• Jie Ren*
• Xidong Feng*
• Bo Liu*
• Xuehai Pan*
• Yao Fu
• Luo Mai
• Yaodong Yang Institute for Artificial Intelligence, Peking University State Key Laboratory of General Artificial Intelligence, Peking University, Beijing, China

Keywords

No keywords are indexed for this paper.