PyTorch: An Imperative Style, High-Performance Deep Learning Library

Adam Paszke,Sam Gross,Francisco Massa,Adam Lerer,James Bradbury,Gregory Chanan,Trevor Killeen,Zeming Lin,Natalia Gimelshein,Luca Antiga,Alban Desmaison,Andreas Köpf,Edward Yang,Zach DeVito,Martin Raison,Alykhan Tejani,Sasank Chilamkurthy,Benoit Steiner,Lu Fang,Junjie Bai,Soumith Chintala
2019-12-04
Abstract:Deep learning frameworks have often focused on either usability or speed, but not both. PyTorch is a machine learning library that shows that these two goals are in fact compatible: it provides an imperative and Pythonic programming style that supports code as a model, makes debugging easy and is consistent with other popular scientific computing libraries, while remaining efficient and supporting hardware accelerators such as GPUs. In this paper, we detail the principles that drove the implementation of PyTorch and how they are reflected in its architecture. We emphasize that every aspect of PyTorch is a regular Python program under the full control of its user. We also explain how the careful and pragmatic implementation of the key components of its runtime enables them to work together to achieve compelling performance. We demonstrate the efficiency of individual subsystems, as well as the overall speed of PyTorch on several common benchmarks.
Machine Learning,Mathematical Software
What problem does this paper attempt to address?
The problem this paper attempts to address is how to improve the usability of deep learning frameworks without sacrificing performance. Specifically, many existing deep learning frameworks such as Caffe, CNTK, TensorFlow, and Theano typically build static data flow graphs to represent the computation process. While this approach helps enhance performance and scalability, it sacrifices ease of use and debugging simplicity. Additionally, some frameworks that adopt dynamic execution (define-by-run) either sacrifice performance (such as Chainer) or use less intuitive languages (such as Torch and DyNet). PyTorch, on the other hand, attempts to combine Python programming style with immediate execution dynamic tensor computation, while also supporting automatic differentiation and GPU acceleration, to achieve an efficient and user-friendly deep learning library. The authors demonstrate PyTorch's design principles and implementation details, explaining how it provides a good user experience while maintaining high performance. The paper also validates PyTorch's performance on various tasks through multiple benchmarks and compares it with several popular frameworks. The results show that PyTorch's performance is comparable to or even better than these frameworks. Furthermore, the high recognition of PyTorch by the research community further proves the effectiveness of its design.