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

TL;DR

Torch.fx is a Python library for capturing and transforming deep learning programs in PyTorch, enabling performance optimization and analysis with a focus on usability for ML practitioners.

Contribution

It introduces a simplified, practical framework for program capture and transformation tailored to typical deep learning workflows in PyTorch.

Findings

Enables workflows previously inaccessible in PyTorch.

Demonstrates practical use cases through case studies.

Optimized for high developer productivity.

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.