Source-linked AI summary
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
TL;DR
Deep learning frameworks often traded usability and flexibility for performance. PyTorch uses imperative, Pythonic execution with automatic differentiation and GPU acceleration, achieving benchmark performance within 17% of the fastest framework while combining usability with careful performance considerations.
Problem
Static dataflow graphs provide advance computation visibility but limit ease of use, debugging, and computational flexibility.
Method
PyTorch combines immediate dynamic tensor execution, automatic differentiation, GPU acceleration, and Pythonic interfaces designed for flexible experimentation.
Results
17%: PyTorch’s performance was within 17% of the fastest framework on all evaluated benchmarks.
Takeaways & Limitations
PyTorch combines usability with performance and became a popular tool in the deep learning research community.
Takeaways & Limitations
The one-pool-per-stream allocator design is susceptible to certain corner cases, although unwanted behavior is rare in practical code.
Abstract
from arXiv · showhide
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.
1 Introduction
PyTorch addresses the tension between deep-learning usability and performance by combining Pythonic dynamic eager execution with automatic differentiation, GPU acceleration, and performance comparable to leading libraries. It is introduced as a Python library designed to make immediate dynamic tensor computation practical without sacrificing speed.
- Motivation: Frameworks such as Caffe [1], CNTK, TensorFlow, and Theano construct static dataflow graphs that represent computations for repeated batch execution.Static graphs provide visibility into the whole computation ahead of time and can theoretically support performance and scalability improvements.
- Related work: Prior dynamic eager-execution frameworks incurred performance costs, as in Chainer, or used less expressive faster languages, as in Torch and DyNet [7].These trade-offs limited the applicability of prior define-by-run approaches.
- Contribution: PyTorch performs immediate execution of dynamic tensor computations with automatic differentiation and GPU acceleration while maintaining performance comparable to the fastest current deep-learning libraries.The paper attributes this combination to careful implementation and design choices.
2 Background
Deep learning builds on four scientific-computing trends: array-based programming, automatic differentiation, open Python ecosystems, and massively parallel hardware. PyTorch combines these trends through GPU-accelerated arrays, automatic differentiation, and integration with Python.
- Domain-specific languages and libraries including APL, MATLAB, R [10], Julia, NumPy [12], Torch [6], Eigen [13], and Lush [14] established tensors and array-based programming as core tools.
- Automatic differentiation automated derivative computation, enabling easier experimentation with machine-learning approaches while retaining efficient gradient-based optimization; autograd popularized it for NumPy arrays.
- The open-source Python ecosystem, including NumPy, SciPy, and Pandas, met researchers' numerical-analysis needs while providing interoperable libraries for broader scientific workflows.
- Commodity GPUs and reusable high-performance kernels such as cuDNN [22] supplied the computing power that enabled frameworks including Caffe [1], Torch7, and TensorFlow to use accelerators.
- PyTorch provides an array-based programming model accelerated by GPUs and made differentiable through automatic differentiation integrated into Python.
3 Design principles
PyTorch’s design balances speed and ease of use through four principles: Pythonic interfaces, researcher-focused usability, pragmatic performance, and simple internal implementations.
- Design principles: Together, these principles weave prior ideas into a design that balances speed and ease of use.The principles guide PyTorch’s overall implementation choices.
- Be Pythonic: PyTorch is Pythonic, using simple, consistent interfaces and integrating naturally with plotting, debugging, and data-processing tools.The goal is to make PyTorch a first-class member of the Python ecosystem with one idiomatic way of doing things.
- Put researchers first: PyTorch hides machine-learning complexity behind intuitive, side-effect-free APIs that make models, data loaders, and optimizers easy and productive to write.This researcher-first design aims to avoid unexpected performance cliffs while keeping common workflows straightforward.
- Provide pragmatic performance: PyTorch accepts implementation complexity to deliver compelling performance without sacrificing simplicity, while giving researchers tools to control execution and improve performance.The design considers a 10% speed tradeoff acceptable for substantially simpler use, but not a 100% tradeoff.
- Worse is better: PyTorch favors a simple but slightly incomplete solution over a comprehensive, complex design when that preserves engineering capacity and maintainability.Saved effort can support additional features, adaptation to new situations, and responsiveness to rapid progress in AI.
4 Usability centric design
PyTorch prioritizes an imperative, Pythonic design in which models, optimizers, and data loaders are ordinary programs, making complex architectures and training techniques easy to implement and debug. Its interchangeable, extensible components also support interoperability with Python libraries and automatic differentiation through dynamically executed, mutating tensor programs.
- Imperative programming: PyTorch preserves imperative programming so arbitrary neural-network architectures can be written as ordinary Python programs, including stateful layer and model classes.This design supports rapidly evolving networks containing loops and recursive functions while keeping new architectures easy to implement.
- Imperative programming: The same programmatic design applies to optimizers and data loaders, enabling training setups such as GANs with two interacting models, optimizers, and losses.Rigid APIs would struggle with GAN training because the losses depend jointly on the generator and discriminator.
- Imperative programming: Eager execution makes Python debugging and inspection tools work throughout development, while exposing intermediate computations without waiting for compilation.Print statements, standard debuggers, and visualization tools such as matplotlib can be used to understand model behavior and verify results.
- Interoperability and extensibility: PyTorch supports bidirectional data exchange with external libraries and lets users extend automatic differentiation and datasets through custom subclasses.Examples include conversion between NumPy arrays and tensors, custom differentiable functions, and new dataset implementations.
- Automatic differentiation: Automatic differentiation uses operator overloading to build a representation of each executed function and tracks tensor mutations with versioning to ensure safe gradients.This approach handles arbitrary Python programs despite Python’s dynamic behavior and supports differentiation through code that mutates tensors.
5 Performance focused implementation
PyTorch achieves high performance in an imperative Python environment by optimizing its C++ core, asynchronous execution, memory management, and multiprocessing while keeping control flow under user control. These mechanisms support GPU utilization, efficient tensor sharing, and immediate memory release despite Python’s execution and memory constraints.
- Execution strategy: PyTorch optimizes every aspect of execution while allowing users to apply additional optimization strategies, avoiding reliance on static data-flow graphs to overcome Python’s interpreter limitations.The approach addresses the global interpreter lock by optimizing execution directly rather than deferring computation to a custom static-graph interpreter.
- Core implementation: Most of PyTorch is implemented in C++, where libtorch provides tensors, CPU and GPU operators, parallel primitives, and automatic differentiation.Generated Python bindings connect this high-performance core to the Python ecosystem.
- Asynchronous execution: PyTorch separates Python-controlled control flow from tensor data flow, then asynchronously queues GPU operators through CUDA streams to overlap CPU execution with GPU work.This overlap keeps the GPU saturated despite Python’s interpreter overhead; operators can run on either CPU or GPU.
- Memory allocation: Its CUDA allocator caches memory, rounds allocations to 512-byte multiples, and maintains a separate pool per CUDA stream to reduce deallocation bottlenecks and fragmentation.Immediate same-stream reuse is enabled because CPU frees precede corresponding GPU reuse under stream serialization, though the design has corner cases and is less suitable for multiple streams.
- Multiprocessing: torch.multiprocessing replaces inefficient tensor serialization with shared memory, improving performance and enabling parallel programs across independent GPUs with later gradient synchronization.It also transparently shares CUDA tensors, supporting techniques such as Hogwild.
- Tensor memory management: PyTorch uses reference counting to release tensor memory immediately when no references remain, but its memory guarantees depend on language runtimes supporting reference counting or user-defined copy and move behavior.This avoids the extra memory usage caused by deferred garbage collection, while unsupported language bindings require specialized memory management.
6 Evaluation
PyTorch achieves competitive single-machine performance across diverse deep-learning tasks, with throughput within 17% of the fastest framework on every benchmark. Its runtime sustains near-perfect GPU utilization through asynchronous execution, while profiling exposes substantial first-iteration CUDA allocation overhead.
- GPU execution: GPU execution takes around three times longer than CPU scheduling, enabling PyTorch to achieve almost perfect device utilization through asynchronous dataflow execution.The result comes from a representative ResNet-50 training-step timeline measured with the built-in profiler.
- CUDA startup overhead: During the first ResNet-50 iteration, cudaMalloc and cudaFree block the CPU thread for long periods, dramatically slowing execution and lowering utilization relative to subsequent iterations.NVIDIA profiling traced both CUDA runtime activity and launched CUDA kernels.
- Overall performance: Within 17% of the fastest framework on every benchmark, PyTorch matches competitive training performance across six models against CNTK, MXNet, TensorFlow, Chainer, and PaddlePaddle.The comparison attributes the similarity partly to shared use of the same cuDNN and cuBLAS versions.
- Community reception: As a proxy for community reception, the evaluation counts monthly arXiv mentions of PyTorch as a percentage of mentions of common deep-learning frameworks after its January 2017 release.The proxy includes Caffe, Chainer, CNTK, Keras, MXNet, TensorFlow, and Theano; the results are reported in Figure 3.
7 Conclusion and future work
PyTorch combines usability with careful performance considerations, supporting its popularity in deep learning research. Future work targets continued improvements in speed and scalability, notably through the PyTorch JIT, which executes programs outside Python for further optimization.
- Conclusion: PyTorch’s popularity in deep learning research reflects its combination of usability and careful performance considerations.
- Future work: Future work will continue improving PyTorch’s speed and scalability, especially through the PyTorch JIT suite for execution outside Python and further optimization.