How to understand Pytorch Source Code?
As an MLE, I feel not satisfied if I don’t have an understanding of the source code of this awesome DL library. As an SWE, I feel I can learn the system design of this great library. This motivated me to explore and write this article.
Based on [1], the fundamental unit in PyTorch is the Tensor. [1] provides an overview of how they implement Tensors in PyTorch, such that the user can interact with it from the Python shell. Four main questions are explained in [1]:
- How does PyTorch extend the Python interpreter to define a Tensor type that can be manipulated from Python code?
- How does PyTorch wrap the C libraries that actually define the Tensor’s properties and methods?
- How does PyTorch cwrap work to generate code for Tensor methods?
- How does PyTorch’s build system take all of these components to compile and generate a workable application?
However, the above doesn’t touch the autograd system though.
The autograd system is moved into C now and is multi-threaded, so stepping through the python debugger is probably a bit pointless. [3]
Here’s a pointer to very old source code, where all the logic was in Python, in the early days of PyTorch. The logic is largely the same still but moved into C.
https://github.com/pytorch/pytorch/blob/v0.1.1/torch/autograd/function.py
Code it up by yourself
The above two links still cannot give us a comprehensive python code to explore what Pytorch implements the autograd. No worries. Watch out for these videos here. Joel Grus is live coding an Autograd Library.
If you are serious 1: READ this blog
This article and this one [A MUST SEE] can be very helpful if you are serious at reading the source code in order to better understand the system or even trying to contribute to the fantastic PyTorch library. I borrowed a figure from the article. I strongly suggest the reader visit the link and read the article by yourself.
PyTorch paper
It does has a paper to introduce the PyTorch with the title of :
PyTorch: An Imperative Style, High-Performance Deep Learning Library
If you are serious 2: read the source code
The codebase structure can be found here.
Critical part: Automatic Differentiation
- A nice video can be found here.
- A nice blog can be found here.
- A nice lecture slides can be found here.
Some repository about the source code
I find this one might be useful.
Reference
[1] https://pytorch.org/blog/a-tour-of-pytorch-internals-1/
[2] https://pytorch.org/blog/a-tour-of-pytorch-internals-2/
[3] https://discuss.pytorch.org/t/how-to-understand-pytorchs-source-code/7600/2
[4]https://en.wikipedia.org/wiki/Automatic_differentiation#Reverse_accumulation
[5]http://blog.christianperone.com/2018/03/pytorch-internal-architecture-tour/
[7]https://zasdfgbnm.github.io/2020/02/07/PyTorch-JIT-Source-Code-Read-Note-Updated-202002/
[8] https://blog.christianperone.com/2019/02/pydata-montreal-slides-for-the-talk-pytorch-under-the-hood/
[9] https://soumith.ch/posts/2021/02/growing-opensource/
[10] http://blog.ezyang.com/2021/01/pytorch-open-source-process/
