標籤：

如何有效地閱讀PyTorch的源代碼？

01-05

看到知乎上有童鞋說PyTorch源碼非常容易閱讀，那麼該如何有效地閱讀PyTorch的源碼呢？謝謝！

相關問題：MXNet的代碼要怎麼讀？

如何高效的學習 TensorFlow 代碼?

代碼統計情況

統計時間: 2017.08.20 (泄漏從啥時候開始看的了...)

  • c、cpp、py都在一個量級，5/6萬左右
  • cuda 3萬

代碼全部都要過一遍的話，需要了解C、C++、Python、CMake、CUDA，外加相關的有CC++代碼打包成Python可調用的.so庫、Python打包安裝過程、BLAS線性庫等。

代碼閱讀順序推薦

  • 運行 python setup.py build ，生成一遍 (非 install，防止覆蓋已安裝的pytorch)
  • 順著 setup.py build 命令看安裝過程，順著安裝過程看相關實現代碼
  • 順著 __init__.py 看 python 中 import torch 時，怎麼把 CC++ 代碼實現的函數與類載入起來的、python層引入了哪些庫

setup.py 中的主要安裝過程

__init__.py 中的主要內容

代碼中比較核心的是TH、THNN，THC、THCNN是對應的GPU版本。

TH中最核心的又是HTStorage、THTensor，前者實現了底層數據(其實就是一個數組)，後者實現了對底層數據的查看。

這樣設計的好處是，有時看起來不一樣的數據底層是共享的，比如矩陣與矩陣的轉置、二維矩陣與二維矩陣變成一維時的矩陣。

一些模塊詳情

  • TH: 核心就THStorage與THTensor，輔助有THGeneral、THHalf、THAllocator、THSize
  • THS: 稀疏矩陣THSTensor(THLongTensor *indices; THTensor *values;)
  • THNN: C代碼的神經網路，每種操作都前向與後向計算，以Abs為例

void THNN_(Abs_updateOutput)(THNNState *state, THTensor *input, THTensor *output)
void THNN_(Abs_updateGradInput)(THNNState *state, THTensor *input, THTensor *gradOutput, THTensor *gradInput)

  • WITH_CUDA: TH、THS、THNN的GPU版本
  • WITH_NCCL: 單機上多GPU通信，支持MPI編程。單PCIe上性能最好，雖然多PCIe也可運行。
  • build THPP: Torch C++
  • build libshm: 共享管理
  • nanopb: 輕量級encode/decode使用Gooogle的Protocol Buffers協議來定義的結構體，以便在RAM或者代碼空間中傳輸消息
  • pybind11: 輕量級的C++與Python間進行無縫調用
  • build ATen: C++的autograd在往這裡寫，整體功能不太清楚
  • generate_nn_wrappers(): 把C代碼寫好的神經網路(NN、CUNN)，使用C++進行封裝，python能調用
  • C: 對C的代碼進行封裝，自身是C++的代碼，核心為Tensor、autograd、cuda，其它有distributed、jit等；
    • torch/csrc/Module.cpp封裝成torch模塊的方法
    • torch/csrc/generic/Tensor.cpp封裝成FloatTensorBase、IntTensorBase等類
  • DL: python可獲取RTLD_GLOBAL、RTLD_NOW變數，來設置載入庫的方式
  • THNN: torch/csrc/nn/THNN.cpp，NN的封裝(generate_nn_wrappers函數生成)
  • THCUNN: torch/csrc/nn/THCUNN.cpp，CUNN的封裝(generate_nn_wrappers函數生成)
  • from torch._C import *: 載入C實現的所有內容(其實是用C++再封裝了一層)

最近在給PyTorch增加對複數的支持（其實我也不想寫代碼，然而沒人給造輪子啊...之前馬普的一個哥們兒已經改了一些CPU和CUDA的部分，但是問題和沒改的還很多）所以大概在兩三個月前開始讀了從上層的Python到底層C和CUDA的實現。感覺只有一開始在C的部分稍微需要動一下腦子，在向量化那一塊可能還需要學習一下SIMD/SSE的指令，其它地方就都很容易讀。

PyTorch的模塊化感覺很好，讀某一個部分不需要很清楚其它部分的具體實現。不過複數是大部分文件里都需要檢查一下要不要改的...

讀底層代碼之前

你需要簡單了解以下介面（知道在哪裡查文檔就好）

  1. BLAS/LAPACK的介面，這個各個BLAS庫都或多或少有文檔來描述，但是個人認為其中intel mkl的最好。有時候還是要參考一下最早的這個BLAS (Basic Linear Algebra Subprograms) 但是注意PyTorch不是調用的CBLAS的介面，而是直接調用Fortran的函數所以所有的變數都是通過指針傳遞，並且相應的routine後面有一個下劃線_
  2. CUBLAS/Thrust這個是CUDA的BLAS部分會用到。這個文檔直接讀NVIDIA官方的就好了。
  3. 相關的數學知識

CPU張量運算庫TH

這個部分是PyTorch里最底層的部分，但是也是我覺得寫的最好的部分。感覺讀懂這一部分其他底層的代碼就都很容易讀了。它用宏實現了C里的泛型支持，然後用宏+命名實現了簡單的OOP。甚至在管理文件的部分還有virtual table。

這個最早是從Lua版本的Torch遷移過來的，那會兒似乎還沒有完整的SIMD。現在提供了完整的SIMD支持。數據結構部分主要實現了`THStorage`和`THTensor`兩個。這兩個是幹嘛用的在PyTorch的Python文檔里都有寫。其實比起我們物理這邊用的Tensor Network的實現要簡單，因為不用連腿....

`THTensor`不負責存儲，是一種查看存儲對象（THStorage）的方法。感覺實現思路上和MXNet的backend，mshadow里的Tensor有一些類似的地方。不過因為C不支持重載，所以這裡完全沒有懶惰求值，當然用的時候也許會有些麻煩。

這裡比較有特點的地方是用宏實現的泛型。這個方法在底層的庫里有廣泛的使用。等到lib目錄上層的代碼能夠使用THPP這個支持C++模板的封裝以後就不會見到了。（THPP和facebook/THPP那個不一樣，後者是一個C++的TH實現）。

TH主要通過宏替換，比如real（這個名字在我改成複數的時候可彆扭了...然後還得小心別和cpp里的std::real一個不小心重名，然後被替換成std::double之類的東西...）真正用來實現的需要泛型的源代碼寫在generic目錄下，然後在外面通過幾個頭文件里定義的宏進行不同變數名稱的替換。

然後還有一個比較難讀的宏應該就是TH_TENSOR_APPLY這個開頭的幾個宏函數了。這幾個比較像是在現在更上層語言里用到的map，或者broadcast。也就是把一段操作應用到每個Tensor的元素上去。為了支持不同個數的Tensor之間的運算，有好幾個不同個數Tensor的。比如TH_TENSOR_APPLY_3就是說三個Tensor里要怎麼做element-wise運算。然後需要注意的是，在宏函數里輸入的操作部分也就是CODE部分，一定不要出現逗號，和內嵌的宏。

然後TensorMath和Conv等部分的函數都和Python函數功能類似，不懂直接看Python文檔就行了。

CPU神經網路庫THNN

這個部分就很簡單了，就是一些Python層的Function對象在底層的定義和實現。不懂的就直接查Python文檔，都有對應的。

自動微分部分

這個是在C++里寫的，然後接了Python。C++里就是實現了三個list用來存Variable啊什麼的...大致思路應該是把使用過forward的函數記下來，然後後面調用backward來算梯度。所以叫Tape-based

剩下的大部分就是Python和一些glue代碼了。Python的文檔很全，更好讀，直接看就行。

不太有空更新，如果有空我會更的...長期待續...不過其實你要是讀了TH部分，基本感覺別的也沒什麼好說的了。

瀉藥

可以先讀一下vision裡面的torchvision.datasets.folder.py這個文件，這是處理圖片數據的，這個部分自定義使用的程度較高，平時自己寫網路也會用到。

接著可以讀一讀nn.module，這裡面幾乎有所有的網路層的定義，然後可以讀一讀激活函數。

剩下的之後再來補充吧。

先從contributing.md開始， 編譯一遍，看看能不能正常編譯，然後跑一遍tests，看看能不能過。最後可以從nn看看module怎麼實現的，要進階就去看看cffi怎麼用。

還是別讀了 要氣死了 有網友報了issue甚至沒人理

你應該從源代碼抽取docstring,這樣生成的數學公式才比較好認

推薦閱讀：

TAG:深度學習DeepLearning | PyTorch |