Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python

Abstract

Modern deep learning frameworks provide imperative, eager execution programming interfaces embedded in Python to provide a productive development experience. However, deep learning practitioners sometimes need to capture and transform program structure for performance optimization, visualization, analysis, and hardware integration. We study the different designs for program capture and transformation used in deep learning. By designing for typical deep learning use cases rather than long tail ones, it is possible to create a simpler framework for program capture and transformation. We apply this principle in torch.fx, a program capture and transformation library for PyTorch written entirely in Python and optimized for high developer productivity by ML practitioners. We present case studies showing how torch.fx enables workflows previously inaccessible in the PyTorch ecosystem.

14 pages, 8 figures, Accepted to MLSys 2022. v2: Added correctness information to evals, clarified 6.2.3, added transform runtime measurement

Cited by in corpus (6)

• Allo: A Programming Model for Composable Accelerator Design
• Oobleck: Resilient Distributed Training of Large Models Using Pipeline Templates
• HAP: SPMD DNN Training on Heterogeneous GPU Clusters with Automated Program Synthesis
• Relax: Composable Abstractions for End-to-End Dynamic Machine Learning
• Performant Automatic Differentiation of Local Coupled Cluster Theories: Response Properties and Ab Initio Molecular Dynamics
• TensorRight: Automated Verification of Tensor Graph Rewrites