Language: 漢語
PyTorch is the the most popular deep learning framework used in both academic and industry. Researchers and engineers use PyTorch to train their models with the convenient and well-defined neural network APIs in eager mode. After training, one has to export the model computation flow as a graph to a framework-agnostic format (e.g., ONNX, TensorRT, etc.) to deploy the trained model for inference with better performance. Thus, we need a mechanism to capture the computation operations happend inside a DNN model.
Graph capturing aims to track and record the PyTorch API invocations used in a model. Since the release of PyTorch 2.0 in 2023, there are two primary machineries to trace a model computation graph: fx.symbolic_trace and Dynamo. In this proposal, I'd like to introduce the usage of both kinds of APIs, and further dissect the underlying working machinary of these two approaches. We also tend to discuss the limitations of both, showing that fx.symbolic_trace is limited by the static symbolic tracing mode and Dynamo is still on its early age.
In this proposal, I'd like to share the following key topics about the graph capture approaches in PyTorch:
- How can we harness fx.symbolic_trace and torch.dynamo.export to capture the DNN operations during the model execution
- How does the symbolic trace mechanism from FX produces the static model graph by substituting the substantial tensors in the graph with "fake" symbolic values
- How Dynamo exploits the custom frame evaluation interface provided by CPython since PEP 523 to dynamically evaluate and dump the PyTorch API invocations during eager execution of DNN models, showing that PyTorch not only treats Python as a frontend language, but also deeply leverages CPython runtime features to extend its flexibility and usability
- The respect limitations of both interfaces