PyTorch的的代碼主要由C10、ATen、torch三大部分組成的。其中:
C10
C10,來自于Caffe Tensor Library的縮寫。這里存放的都是最基礎的Tensor庫的代碼,可以運行在服務端和移動端。PyTorch目前正在將代碼從ATen/core目錄下遷移到C10中。C10的代碼有一些特殊性,體現在這里的代碼除了服務端外還要運行在移動端,因此編譯后的二進制文件大小也很關鍵,因此C10目前存放的都是最核心、精簡的、基礎的Tensor函數和接口。
C10目前最具代表性的一個class就是TensorImpl了,它實現了Tensor的最基礎框架。繼承者和使用者有:
Variable的Variable::Impl
SparseTensorImpl
detail::make_tensor(storage_impl, CUDATensorId(), false)
Tensor(c10::intrusive_ptr<TensorImpl, UndefinedTensorImpl> tensor_impl)
c10::make_intrusive<at::TensorImpl, at::UndefinedTensorImpl>
值得一提的是,C10中還使用/修改了來自llvm的SmallVector,在vector元素比較少的時候用以代替std::vector,用以提升性能;
ATen
ATen,來自于 A TENsor library for C++11的縮寫;PyTorch的C++ tensor library。ATen部分有大量的代碼是來聲明和定義Tensor運算相關的邏輯的,除此之外,PyTorch還使用了aten/src/ATen/gen.py來動態生成一些ATen相關的代碼。ATen基于C10,Gemfield本文討論的正是這部分;
torch._VF
torch._VF is a subset of those functions are mapped to ATen functions in torch/jit/_builtins.py.
例如如果沿著torch.nn.dropout去找,會發現找到torch._VF這個庫,IDE無法繼續navigate:
這時候直接去github的pytorch源碼repo(https://github.com/pytorch/pytorch/tree/b3bb234e16780ea3f3d749d2c8c156c9245eb797)搜dropout_就能很快看到在cpu上cpp的實現:
加上cuda的限制就是gpu的實現:
另外在pytorch的老舊版本,像dropout這種由python實現的算子在pytorch v1.0.4版本的https://github.com/pytorch/pytorch/blob/v0.1.4/torch/nn/functions/dropout.py 這個目錄下還能找到,但是github并沒有提供歷史tag的搜索功能,不如就去https://github.com/colesbury/pytorch-old/tree/master/torch/nn/functions 這個古早目錄下面找吧
Torch
Torch,部分代碼仍然在使用以前的快要進入歷史博物館的Torch開源項目,比如具有下面這些文件名格式的文件:
TH* = TorcH
THC* = TorcH Cuda
THCS* = TorcH Cuda Sparse (now defunct)
THCUNN* = TorcH CUda Neural Network (see cunn)
THD* = TorcH Distributed
THNN* = TorcH Neural Network
THS* = TorcH Sparse (now defunct)
THP* = TorcH Python
PyTorch會使用tools/setup_helpers/generate_code.py來動態生成Torch層面相關的一些代碼,這部分動態生成的邏輯將不在本文闡述,你可以關注Gemfield專欄的后續文章。
部分轉載自
- https://zhuanlan.zhihu.com/p/55966063