Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python
This provides a practical tool for deep learning practitioners to optimize, analyze, and integrate hardware in PyTorch, though it is incremental as it builds on existing program transformation concepts.
The paper tackles the need for program capture and transformation in deep learning frameworks by introducing torch.fx, a library for PyTorch that simplifies these tasks for typical use cases, enabling previously inaccessible workflows in the PyTorch ecosystem.
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.