PyTorch 源碼學習：閱讀經驗代碼結構

文章目錄

[1.0.1] Oldpan：Pytorch源碼編譯簡明指南
- 核心文件夾
- third_party
- tools
[1.7.1] 孫港丶丶：查看 Pytorch 源碼
[1.10.2] jhang的回答?：如何閱讀pytorch框架的源碼？
[2.0.0]吾乃阿爾法：一點 PyTorch 源碼閱讀心得
[2.1.1]kiddyjinjin：pytorch 源碼解讀進階版 - 總述
[unknown] 小飛的回答?：如何閱讀pytorch框架的源碼？
[unknown] 小飛?：PyTorch源碼學習系列 - 1.初識
- 我理解的PyTorch架構
- PyTorch目錄結構
[unknown] clay001?：pytorch源碼閱讀
待更新……

[1.0.1] Oldpan：Pytorch源碼編譯簡明指南

以下是Pytorch源碼包展開的目錄結構(只展示了主要的一些文件夾)，其中主要的源碼都在以下展示的文件夾中：

《Pytorch源碼編譯簡明指南》

其中使用紅箭頭標注的就是幾個比較重要的核心庫。下面簡單介紹一下：

核心文件夾

核心文件夾主要是c10、aten、torch、caffe2.

為什么將c10放到最前面呢？

因為官方已經表明c10目錄是最重要的源代碼文件夾，也就是幾乎所有的源代碼都與這里的代碼有關系，比如我們的類型定義，Pytorch最重要的Tensor的內存分配方式等等，都在這個文件夾中，官方也說到了，之后會慢慢將Aten中的代碼移至這個文件夾，也就是說這個文件夾將包含Pytorch中最核心的代碼。

而Aten文件夾則包含了一些實現了Tensor的底層(和c10類似)，也包括了很多的層前向代碼和后向實現的代碼(例如卷積層的前向和后向操作代碼)，包括CPU和GPU端，總之都是C++的核心操作代碼。

torch文件夾也同樣重要，其中主要包含了一些稍微高層些的操作函數，例如torch.ones等，有C++和Python端，也包括了Python核心代碼和包裝代碼，如果我們使用python版Pytorch的話，與這些代碼接觸就比較密切了。

而Caffe2則不用多說，caffe2則主要針對移動端設計了很多優化后的運算代碼，模型融合、模型量化等等的代碼，其后端有QNNPACK等一些針對移動端的底層運算庫(有開發人員說GLOW也在caffe2后端考慮之內)。

third_party

Pytorch畢竟是大型的深度學習庫，所以需要的依賴庫也是有很多的，其中有很多我們耳熟能詳的數值計算庫(eigen、gemmlowp)、模型轉換庫(onnx、onnx-tensorrt)、并行訓練庫(gloo、nccl)、自家的底層端實現庫(QNNPACK)以及綁定python端的pybind11等一系列所依賴的庫。

當然還有很多庫這里就不一一介紹了，總之，我們在編譯的時候，Pytorch的編譯代碼會根據我們的設置在編譯的時候，自動判斷當前系統中是否存在需要的第三方庫。如果不存在則使用這里的第三方庫(直接編譯并使用第三方庫的diamante)，這也是為什么我們需要執行git submodule update --init --recursive來下載所依賴第三庫源碼的原因。

tools

tools這個文件夾中的內容到底是做什么的，簡單看一下官方的介紹：

This folder contains a number of scripts which are used as part of the PyTorch build process. This directory also doubles as a Python module hierarchy. 此文件夾包含許多用作PyTorch構建過程一部分的腳本。該目錄還兼作Python模塊層次結構。

其中包含了一些腳本生成代碼工具(利用python)、用于編譯一些組件的腳本和代碼，還有一些開發人員需要的工具、以及AMD顯卡幫助編譯代碼和一些特殊情況需要使用的工具等。在我們編譯Pytorch源碼的過程中會使用到這個文件夾中的代碼。

有一點需要說明，那就是Pytorch利用了很多的代碼生成，例如操作層函數的頭文件NativeFunction.h等，所以tools中的代碼生成腳本還是比較重要的。

提一個可能會使用到的腳本build_pytorch_libs.sh，這個腳本是用來編譯libtorch庫的，libtorch就是不需要python包裝的使用C++的Pytorch庫，方便于部署階段使用。

關于tools中的文件就不具體介紹了，大家可以看一下其中的readme。

其他的文件夾就不多說了，相對上面的來說并不是很重要。

[1.7.1] 孫港丶丶：查看 Pytorch 源碼

在這里插入圖片描述

“torch/” 中的代碼文件一般是pytorch python類的接口信息，其內容可以直接在編輯器中通過查看定義得到，但其只包括了接口的參數信息和注釋，注釋與官方文檔中收到的內容相同；
“torch/csrc/” 中含有python和c++銜接的代碼，很多為編譯時生成；
“aten/src/ATen/” 中包括了torch中許多操作的C++或C實現，是我們查看pytorch許多函數實現需要關注的區域；
“c10/” 中為torch最基本特性實現的部分，如張量的定義、數據類型的定義等。

因此，為查看torch函數的具體實現，需要查看"aten/src/ATen/" 部分的內核代碼。

其中"native/“為C++風格的實現，”/TH"等主要為C風格的實現，Edwards 說開發人員在努力將"/TH"等中的操作遷移到"native/"中，至少在我使用的pytorch1.7.1中，"native/"覆蓋了大部分的張量操作，如卷積、激活和損失等等。

如何查看"aten/src/ATen/native"中的函數操作？

首先可以在"aten/src/ATen/native/native_functions.yaml"這個注冊文件中查看函數的框架，可以確定這個函數是否在"native/"中實現，這里的框架主要是用于代碼生成。
直接搜索大法，搜索文件夾中所有文件的內容，找到你想要的函數。\

[1.10.2] jhang的回答?：如何閱讀pytorch框架的源碼？

要想閱讀 pytorch 框架的源碼，最起碼得先了解框架的調用棧，以 pytorch v1.10.2的 torch.nn.conv2d 為例，介紹一下python api 接口到底層 c++ 接口的調用依賴關系：

Python API:

torch.nn.Conv2d:
---> class Conv2d, (torch/nn/modules/conv.py)
---> F.conv2d, (torch/nn/functional.py)
---> def conv2d(input: Tensor, ...) -> Tensor: ..., (torch/_C/_VariableFunctions.pyi)

C++ API:

---> at::Tensor conv2d(...), (aten/src/ATen/native/Convolution.cpp: 579)
---> at::convolution(...), (aten/src/ATen/native/Convolution.cpp: 586)
---> at::Tensor convolution(...), (aten/src/ATen/native/Convolution.cpp：743)
---> at::_convolution(...), (aten/src/ATen/native/Convolution.cpp：754)
---> at::Tensor _convolution(...), (aten/src/ATen/native/Convolution.cpp：769)
---> at::cudnn_convolution(...), (aten/src/ATen/native/Convolution.cpp：860)
---> Tensor cudnn_convolution(...), (aten/src/ATen/native/cudnn/ConvShared.cpp：265)
---> cudnn_convolution_forward(...), (aten/src/ATen/native/cudnn/ConvShared.cpp：273)
---> Tensor cudnn_convolution_forward(...), (aten/src/ATen/native/cudnn/ConvShared.cpp：221)
---> raw_cudnn_convolution_forward_out(...), (aten/src/ATen/native/cudnn/ConvShared.cpp：258)
---> void raw_cudnn_convolution_forward_out(...), (aten/src/ATen/native/cudnn/Conv_v7.cpp：669)
---> split_batch_dim_to_32bit_out(...,raw_cudnn_convolution_forward_out_32bit), (aten/src/ATen/native/cudnn/Conv_v7.cpp：673)
---> void raw_cudnn_convolution_forward_out_32bit(...), (aten/src/ATen/native/cudnn/Conv_v7.cpp：625)
---> cudnnConvolutionForward(...), (aten/src/ATen/native/cudnn/Conv_v7.cpp：658)

python 端接口比較容易查找，對于 c++ api 接口的查找給與一些建議：

常用 op c++ 接口都位于aten/src/ATen/native目錄，而且該目錄下有個native_functions.yaml, 記錄了c++接口的調用依賴關系
kernel 名搜索：通常我們也會通過 nvprof 查看調用 kernel 名逆向查找 c++調用接口
關鍵字搜索：以conv2d為例, 我們也會去搜索 Tensor conv2d 關鍵字來查找，方便將 python 端的最后一個接口與 c++ 端的第一個接口對應起來；
先驗知識搜索，還需要擁有一些cudnn api的先驗知識，比如我知道conv2d一定會調用 cudnnConvolutionForward 等接口, 然后搜該接口所在位置，然后逆向搜索 c++ api 接口即可
c++ 接口的搜索一般只需要找到離 cuda kernel 或 cudnn/cublas API 最近的位置即可(所以方法 2，4是我們最常用的手段)，從 python 接口到 c++ 接口的中間依賴關系沒必要深究，因為其中做了很多封裝，經常依賴 vscode 跳轉時會斷了依賴關系；
推薦使用vscode開發工具做好配置，方便c++/python 接口的跳轉，有助于理解調用過程

[2.0.0]吾乃阿爾法：一點 PyTorch 源碼閱讀心得

PyTorch 的代碼對初學者來說并不好讀，原因在于：

有很多 C++ 代碼是在編譯過程中生成的；
PyTorch 前端充分利用了 Python 的動態特性，比如 Python 函數可能經過 N 個裝飾器修飾；
混合 Python/C++ 調用，比如從 Python 調 C++ 調 Python 再調 C++ 在 PyTorch 中是比較常見的事情；

以下是我讀 PyTorch 源代碼的一些經驗，僅供參考：

明確目標: PyTorch代碼很多，建議每次只讀一個專題或者從一個特定的問題出發，比如 PyTorch AMP 是怎么實現的；
把握全局: 一上來直接讀代碼會有很多障礙，很多術語不明所以，先通讀這個專題下官方的教程、文檔，以及一些寫的好的第三方博客，確保自己對這個專題下的內容有初步的認知。以下是一些初步了解 PyTorch 特定專題內容的比較好的資源:
- PyTorch 教程
- PyTorch 文檔
- PyTorch 博客
- PyTorch 案例
- PyTorch 論壇
- PyTorch Youtube 頻道
Debug Build: 一定要 build debug 版的 PyTorch，并保留在編譯過程中生成的源代碼，否則很難找到 PyTorch 函數調用棧中一些函數的來源；

此處省略細節，具體內容見原文
靜態讀代碼: 有了完整的 PyTorch 源代碼之后就可以開始讀了，網上有很多 VSCode 教程，設置好 VSCode 的 Python 和 C++ 插件，方便在函數之間跳轉，可以解決一大部分的函數跳轉；
動態讀代碼: 靜態讀代碼的問題是常常搞不清函數的執行流程，此時在運行過程中動態讀執行過的代碼就很有幫助，善用 gdb 和 pdb 可以有效輔助讀源代碼。

此處省略細節，具體內容見原文
充分利用源代碼中的日志、debug 選項、測試用例: 很多 PyTorch 模塊都包含了豐富的日志和 debug 開關，這些日志和用于調試的消息可以幫助我們理解 PyTorch 的執行流程。除此之外，PyTorch 中包含了大量的測試用例，如果單純看源代碼無法理解程序的邏輯，看看其對應的測試用例可以幫助我們理解程序在做什么。

此處省略細節，具體內容見原文
及時求助: 如果經過上面的流程還無法了解某些代碼的邏輯，要及時向社區求助，避免浪費過多時間。
學什么: 明確源代碼中哪些東西值得我們學習和借鑒，讀源代碼時要特別注意這些方面，比如:
- 特定模塊/功能的實現原理;
- 用到的算法;
- 一些 coding 技巧;
知行合一: You can’t understand it until you change it. 讀源代碼不是最終目的，充分利用從代碼中獲取的認知才是。有效的輸出可以加深我們對代碼的理解，一些可以參考的輸出方式:
- 寫一篇源碼剖析的博客;
- 簡化自己對源代碼的認識，分享給其他人;
- 修改源代碼，改進或添加一些功能，給 PyTorch 提交 PR;
- 親手實現相同功能，寫精簡版的代碼復現核心邏輯;

每一個讀源代碼的人都是不甘平凡的人，祝大家在這個“痛并快樂著”的過程中成長得更快、更多。

[2.1.1]kiddyjinjin：pytorch 源碼解讀進階版 - 總述

pytorch官方給出的參考代碼結構：pytorch/pytorch

源碼解讀的代碼版本基于v2.1.1 版本：GitHub - pytorch/pytorch at v2.1.1

其中比較核心的4個目錄簡單介紹如下：

torch：python 代碼部分， the frontend, 我們在代碼中引入并使用的Python 模塊(modules)，基本都在這里
torch/csrc: python bindings, 這部分C++代碼實現了所謂的PyTorch前端(the frontend of PyTorch)。具體來說，這一部分主要橋接了Python邏輯的C++的實現，和一些PyTorch中非常重要的部分，比如自動微分引擎(autograd engine)和JIT編譯器(JIT compiler)。
aten/src/ATen: 是“A Tensor Library”的縮寫，是一個C++庫實現了Tensor的各種operations。ATen 內對operators的實現分成兩類，一種是現代的C++實現版本（native），另一種是老舊的C實現版本（legacy）。
c10: 是一個來自于Caffe2 和 ATen的雙關語(Caffe 10)，其中包含了PyTorch的核心抽象，比如 Tensor、Device、Allocator、Storage 等數據結構的實際實現部分。

進一步對c10/core 的代碼結構進一步介紹如下：

Allocator.h/cpp: 提供memory管理的抽象類定義和實現。
CPUAllocator.h/cpp: set/get CPU Allocator, 指定DefaultCPUAllocator。
Device.h/cpp: 提供底層硬件的抽象類定義和實現。
- struct C10_API Device final 表示張量所在的計算設備, 設備由類型唯一標識（cpu or cuda gpu， cuda 的話有index）
DeviceGuard.h : RAII guard
DeviceType.h 定義了21 種 DeviceType，包括我們常見的CPU，CUDA，XLA，HIP 等，應該會是逐漸增加的狀態。
DispatchKey.h 定義了40+種 dispatcherkey set，是個 uint8_t，越高位優先級越高，用來給op路由，當調用一個op的時候，pytorch中的dispatch 機制會結合 DispatchKey 以及當時的場景（在什么硬件上？是推理還是訓練？）分配到特定的執行算子上，這個機制保證了pytorch可以非常靈活的在不同的硬件和場景下切換。后面會詳細講講，比較有意思。
DispatchKeySet.h
TensorOptions.h/cpp: 提供創建張量時的選項（如設備類型、數據類型等）的定義和相關操作。
TensorImpl.h/cpp: 提供張量的具體實現。
其他：
Backend.h （老舊版本的 Device.h/cpp， layout 信息），提供多種新舊互相轉換接口

[unknown] 小飛的回答?：如何閱讀pytorch框架的源碼？

分享下我個人學習經驗

首先自己一定要編譯源碼，學會用GDB去調試代碼。任何源碼學習的第一步永遠是將源碼編譯成功并運行起來。其次切記不要去看 setup.py 文件，不要去看 torch/__init__.py 文件，不要去看C++編譯配置文件。如果你沒有豐富的相關工程開發經驗，從你打開這些文件起，恭喜你，你已經被成功勸退了。初學者就應該老老實實先學習項目的設計思想和架構，忽略工程配置。
了解下Python語言如何與C/C++語言綁定。PyTorch底層核心代碼都是用C/C++編寫，如果你僅僅想學習的是Python部分的源碼，那本回答就到此結束了，根據你python里的模塊用到哪學到哪就行了。如果你想深入學習C/C++源碼，這是一道必須邁過的坎。不理解的話你對PyTorch的認識永遠是Python和C/C++兩個孤零零的世界，無法聯系起來。
先學習Tensor。Tensor是貫穿整個PyTorch最基本的數據結構，無論何時你都會遇到它。
了解算子。通過最簡單的Tensor.Add 方法學習下Tensor函數是如何設計的。PyTorch里操作Tensor的函數被稱之為算子，算子包含前向和后向計算。這里不要求完全弄懂算子，只需要了解相關流程就行。
學習Autograd。先了解計算圖的概念，學習下PyTorch是如何創建和維護計算圖。如何反向傳播自動求解梯度更新權重。

學到這里你基本上對PyTorch C/C++底層架構有所了解，下面就是根據自己喜愛的方向去學習。

如果你想研究深度學習模型，去看Python torch.nn包里面的python源碼就行。
如果你想研究算子實現，還記得我們前面學的Add方法了，找到你想了解的算子去深入，如果不關心速度，看下CPU代碼就行，如果想學習如何優化就去了解cuda編程。
如果想學習分布式，就去研究下分布式相關的code。
其實從上面的流程看，只要你前5步都完成了，你自然而然就知道該如何去學習PyTorch框架源碼。現在PyTorch源碼很大，想在有限時間內一下子全看完不現實，必然要結合自己的實際情況有所側重。

[unknown] 小飛?：PyTorch源碼學習系列 - 1.初識

PyTorch本身是一個Python擴展包，按照官方說法它主要具有以下兩種特色：

支持GPU加速的張量（Tensor）計算
在一個類似磁帶（前向和反向）的梯度自動計算（Autograd）系統上搭建深度神經網絡

Tensor其實本質上就是一個多維數組。在數學上單個數據我們稱之為標量，一維數據我們稱之為向量，二維數據我們稱之為矩陣。

如果PyTorch僅是支持GPU加速的Tensor計算框架，那它也就是NumPy的替代品而已。其最核心的功能就是Autograd系統，目前深度學習就是基于梯度反向傳播理論來達到網絡的自我訓練。PyTorch的Autograd系統是創建了一個動態的計算圖，用戶只需要關注前向計算網絡搭建，PyTorch會自動根據動態計算圖去反向計算梯度并更新網絡權重。

在設計之初PyTorch并沒打算僅成為一個綁定C++框架的Python包，它緊密地將C++框架集成到python中。你可以在使用PyTorch的同時也結合使用NumPy/SciPy/scikit-learn這些優秀的科學計算包，同時你也可以用Python編寫你自己的神經網絡層（PyTorch神經網絡相關的代碼基本上都是Python編寫）。

正如前面所說，PyTorch的計算圖是動態的，當你寫下一行Python代碼的時候你就可以直接運行得到結果。對用戶來說，PyTorch的世界就是一個簡單的單線程同步非阻塞世界，你可以得到實時反饋，這對于剛開始接觸深度學習的新手來說是非常貼心的功能，可以幫忙新手快速入門。

我理解的PyTorch架構

根據自己的理解簡單畫了下PyTorch的架構圖，隱藏了部分細節。本系列也將主要圍繞著這張架構圖去學習PyTorch的具體實現。

在這里插入圖片描述

一共將PyTorch分成了四層，分別是

應用層（Python）。這應該是大家最熟悉的層，主要涉及到張量，Autograd以及神經網絡。該層所有的源碼都是由Python編寫，這也符合前面所說的PyTorch設計思想-——將C++框架集成到Python里
實現接口層（C++）。該層的主要功能我認為有兩個：
- Python 擴展。通過Python提供的C API將Python應用層與C++實現層綁定起來，使用戶在享受Python語言提供的便捷優勢時也可以同時享受到C++語言提供的性能優勢
- Autograd系統實現。 PyTorch并沒有在實現層中實現Autograd系統。在此層中PyTorch定義了動態有向圖的基本組件Node和Edge，以及在此基礎上封裝了Function類和Engine類來實現Autograd
實現層（C++）。該層是PyTorch的核心層，定義了PyTorch運行過程中的核心庫，包括Tensor的具體實現，算子實現（前向與后向運算）以及動態調度系統（Tensor的布局，硬件設備，數據類型）。Storage類主要是針對不同硬件數據存儲的一種抽象。
硬件接口層。該層主要是硬件廠商基于自家硬件推出的運算接口。

PyTorch目錄結構

PyTorch的源碼托管于GitHub平臺，其目前的代碼量已經非常巨大。新手第一次接觸的時候往往會因此被勸退，但其實里面很多文件和我們學習PyTorch源碼并沒有太多的直接關系，所以我們第一步就是要理清目錄結構，專注于我們需要學習的內容。

torch：我們“import torch”后最熟悉的PyTorch庫。所有非csrc文件夾下的內容都是標準的Python模塊，對應我們架構圖中的應用層
- csrc：該目錄下都是C++源碼，Python綁定C++的相關code都在這個目錄里面，同時也包含了對PyTorch核心庫的一些封裝，對應我們架構圖中的實現接口層
- csrc/autograd：梯度自動計算系統的C++實現
- autograd：梯度自動計算系統的Python前端源碼，包含torch中支持的所有自動求導算子
- nn：建立在autograd系統上的神經網絡庫，包含了深度學習中常用的一些基礎神經網絡層。
- optim：機器學習中用到的優化算法庫
aten：“a tensor library”的縮寫，對應我們結構圖中的實現層。從名字上也能知道，這個庫設計之初主要是為Tensor服務。因為在實現接口層下面，所以這里的Tensor并不支持autograd
- src/Aten/core：aten的核心基礎庫。目前這個庫里面的代碼正在逐漸地遷移到c10目錄下面
- src/Aten/native：PyTorch的算子庫，這個目錄下面的算子都是CPU的算子。對于一些專門CPU指令優化的算子會在子目錄里面
- src/Aten/native/cuda：cuda算子實現
c10：“caffe2 aten”的縮寫，PyTorch的核心庫，支持服務端和移動端。
tools：PyTorch中很多相似源碼都是腳本通過模板自動生成的，這個文件夾下面就放著自動生成代碼的腳本

[unknown] clay001?：pytorch源碼閱讀

pytorch代碼主要由C10（Caffe Tensor Library，最基礎的Tenor庫代碼），ATen（A Tensor library for C++11，基于C10），torch三大部分組成。

torch.nn中包含各種神經網絡層，激活函數，損失函數的類，使用時需要先創建對象。

torch.nn.functional中的接口可以直接調用函數而不用創建對象。

算子配置文件：pytorch/aten/src/Aten/native/native_functions.yaml中有關于各個算子的說明，同級目錄下有這些算子的實現。每個算子有一些字段：func，variants，dispatch。

func字段：表示算子的名稱和輸入輸出類型
variants字段：表示高級方法的使用方式，
- function表示可以用torch.xxx()的方式調用，
- method表示在tensor a上，用a.xxx()的方式調用
dispatch字段：分發的設備類型，CPU，CUDA等等