Back to the era of deep learning with Tensorflow, researchers and engineers were more or less complaining the difficulties of using Tensorflow. And the main pain poit is you need to build the whole computation graph (model definition, optimizer, etc) before you start to train your model. Then PyTorch came into the game with a much more intuitive eager mode, where you can run your model line by line anytime.

However, what is the cost? Answer to this boils down to efficiency. Since eager mode needs to dynamically build computation graph and does not apply potential computation optimizations (like fusing operations), efficiency in both training and inference (including computation and memory access) is largely impacted. When moving to deployment, it can cause significant difference.

This post is aimed to understand how the graph mode in PyTorch, and key techniques for enabling graph mode in PyTorch.

To start with, the following lists the development of PyTorch graph mode:

mode 2018 2021 2023
module level tracing torch.jit.trace torch.fx.symbolic_trace  
bytecode level tracing torch.jit.script   torch.dynamo