PyTorch源碼淺析（目錄）

04-13

引子

17年６月開始由於在我們的問題求解過程中需要使用到複數，而唯一支持複數的框架tensorflow也並不是所有的操作都支持，此外對於複數這個部分優化的有限，所以著手開始嘗試修改幾個流行的框架。我分別在MXNet的後端mshadow，PyTorch的倉庫下提了issue試圖詢問開發者是否有相關計劃，不過遺憾的是好像由於目前複數在實際工作中的應用並不是很多，複數支持並非是優先順序較高的特性。

首先感謝一下這期間中國科學技術大學韓永健老師，中科院物理所王磊老師，復旦大學李曉鵬老師組的支持（沒阻止我不務正業，但是沒輪子用很難受啊，捂臉）。

雖然從底層代碼的角度來看，個人更加喜歡MXNet的C++代碼，和PyTorch的幾個C語言庫。讀過一些很有收穫。但是只有PyTorch的開發者鳥我了，所以就開始改底層有些不太好改的PyTorch（相比MXNet而言）。相比之下當時還在v0.1版本的PyTorch底層就混亂了許多，這也導致從底層支持複數並不是那麼容易。第一個完成複數支持的倉庫是在Github:PhillippPelz修改的版本之上做的，但是他的做法似乎PyTorch的團隊並不是很喜歡，他直接修改了幾個後端的C/CUDA C庫，並且導致一些命名上的混亂。後來我在之前的基礎上將複數單獨拿出來，交了第一個PR，大約９k行的C，用以提供一個支持複數的CPU後端，但是也許是因為太多了，沒人review也沒人merge。期間有人聯繫過我向我要源碼，相關指導等等，可見複數並非我一個人所需。如果不儘早merge，由於PyTorch的更新速度比較快（三個月一個release）我一個人很難跟上它的進度，因為整體工程量不小，等我修改完又需要費很大力氣去跟進其它新特性。所以更希望是先能merge　CPU然後再支持GPU。（GPU上用float2來模擬複數的時候map/reduce之類的地方一些原來能用trick變得很不好寫）跨過幾個版本可能就又需要很多修改。所以後來我就放棄獨立支持複數的計划了，改去使用Julia，並且在Julia上搭建自己所需。但是在修改PyTorch期間，為了方便我記錄了一些筆記來闡述源碼的構成和作用，以及我猜測的作者的設計思路，現在慢慢整理分享出來以饗讀者。

整體結構

目前我本地版本是commit d0cabbd

PyTorch是一個以Python為前端的框架，在藉助Python這個動態語言本身的特性基礎上使用動態圖來定義模型，可以說PyTorch很好地藉助了Python編譯器的一些特性（例如垃圾回收，語言本身的動態性等），這使得用戶在使用時會覺得好像比tensorflow這樣使用靜態圖的框架容易。

結構上大約有兩部分，一部分是Python，另外一部分是C/C++寫的後端。Python的部分比較容易讀，如果後面有時間再介紹。C/C++後端的多維數組庫在v0.4開始用aten(a tensor library)庫作為上層封裝再通過torch/csrc 中的部分膠水代碼接入Python（但慢慢在改為pybind11）。自動微分使用C++完成，後端採用的是一個類似於錄像帶的機制來記錄相關操作。新加入的jit部分還有詳細讀，不介紹。用`tree`看一下大約是這樣的。後端核心代碼主要在 lib 里。

├── autograd├── backends├── contrib...├── csrc│ ├── allocators.cpp│ ├── allocators.h│ ├── assertions.cpp│ ├── assertions.h│ ├── autograd│ ├── byte_order.cpp│ ├── byte_order.h│ ├── copy_utils.h│ ├── cuda│ ├── DataLoader.cpp│ ├── distributed│ ├── dl.c│ ├── DynamicTypes.cpp│ ├── DynamicTypes.h│ ├── Exceptions.cpp│ ├── Exceptions.h│ ├── expand_utils.cpp│ ├── expand_utils.h│ ├── Generator.cpp│ ├── Generator.h│ ├── generic│ ├── jit│ ├── Module.cpp│ ├── Module.h│ ├── ModuleSparse.cpp│ ├── nn│ ├── nnpack│ ├── nvrtc.cpp│ ├── onnx│ ├── PtrWrapper.cpp│ ├── PtrWrapper.h│ ├── PythonTypes.h│ ├── README.md│ ├── serialization.cpp│ ├── serialization.h│ ├── Size.cpp│ ├── Size.h│ ├── Storage.cpp│ ├── Storage.h│ ├── Tensor.cpp│ ├── Tensor.h│ ├── THP_API.h│ ├── THP_export.h│ ├── THP.h│ ├── Types.h│ ├── utils│ ├── utils.cpp│ └── utils.h├── cuda...├── distributed...├── distributions.py├── for_onnx...├── jit├── legacy├── lib│ ├── ATen│ ├── build│ ├── build_libs.bat│ ├── build_libs.sh│ ├── gloo│ ├── include│ ├── libshm│ ├── libshm_windows│ ├── libTH.so.1│ ├── nanopb│ ├── nccl│ ├── pybind11│ ├── README.md│ ├── TH│ ├── THC│ ├── THCUNN.h│ ├── THD│ ├── THNN.h│ └── tmp_install├── multiprocessing├── nn├── onnx├── optim├── random.py├── serialization.py...├── _thnn├── _torch_docs.py├── utils└── _utils.py53 directories, 110 files