Tensor
自從張量（Tensor）計算這個概念出現后，神經網路的演算法就可以看作是一系列的張量計算，所謂的張量，它原本是個數學概念，表示各種向量或者數值之間的關系，PyTorch的張量（torch.Tensor）表示的是N維矩陣與一維陣列的關系，
torch.Tensor的使用方法和numpy很相似（https://pytorch.org/...tensor-tutorial-py），兩者唯一的區別在于torch.Tensor可以使用GPU來計算，這就比用CPU的numpy要快很多，
張量計算的種類有很多，比如加法、乘法、矩陣相乘、矩陣轉置等，這些計算被稱為算子（Operator），它們是PyTorch的核心組件，
算子的backend一般是C/C++的拓展程式，PyTorch的backend是稱為"ATen"的C/C++庫，ATen是"A Tensor"的縮寫，
Operator
PyTorch所有的Operator都定義在Declarations.cwrap和native_functions.yaml這兩個檔案中，前者定義了從Torch那繼承來的legacy operator（aten/src/TH），后者定義的是native operator，是PyTorch的operator，
相比于用C++開發的native code，legacy code是在PyTorch編譯時由gen.py根據Declarations.cwrap的內容動態生成的，因此，如果你想要trace這些code，需要先編譯PyTorch，
legacy code的開發要比native code復雜得多，如果可以的話，建議你盡量避開它們，
MatMul
本文會以矩陣相乘--torch.matmul()為例來分析PyTorch算子的作業流程，
我在深入淺出全連接層（fully connected layer）中有講在GPU層面是如何進行矩陣相乘的，Nvidia、AMD等公司提供了優化好的線性代數計算庫--cuBLAS/rocBLAS/openBLAS，PyTorch只需要呼叫它們的API即可，
Figure 1是torch.matmul()在ATen中的function flow，可以看到，這個flow可不短，這主要是因為不同型別的tensor（2d or Nd, batched gemm or not，with or without bias，cuda or cpu）的操作也不盡相同，
at::matmul()主要負責將Tensor轉換成cuBLAS需要的格式，前面說過，Tensor可以是N維矩陣，如果tensor A是3d矩陣，tensor B是2d矩陣，就需要先將3d轉成2d；如果它們都是>=3d的矩陣，就要考慮batched matmul的情況；如果bias=True，后續就應該交給at::addmm()來處理；總之，matmul要考慮的事情比想象中要多，
除此之外，不同的dtype、device和layout需要呼叫不同的操作函式，這部分作業交由c10::dispatcher來完成，
Dispatcher
dispatcher主要用于動態呼叫dtype、device以及layout等方法函式，用過numpy的都知道，np.array()的資料型別有：float32, float16，int8，int32，.... 如果你了解C++就會知道，這類程式最適合用模板（template）來實作，
很遺憾，由于ATen有一部分operator是用C語言寫的（從Torch繼承過來），不支持模板功能，因此，就需要dispatcher這樣的動態調度器，
類似地，PyTorch的tensor不僅可以運行在GPU上，還可以跑在CPU、mkldnn和xla等設備，Figure 1中的dispatcher4就根據tensor的device呼叫了mm的GPU實作，
layout是指tensor中元素的排布，一般來說，矩陣的排布都是緊湊型的，也就是strided layout，而那些有著大量0的稀疏矩陣，相應地就是sparse layout，
Figure 2是strided layout的演示實體，這里創建了一個2行2列的矩陣a，它的資料實際存放在一維陣列（a.storage）里，2行2列只是這個陣列的視圖，
stride充當了從陣列到視圖的橋梁，比如，要列印第2行第2列的元素時，可以通過公式：\(1 * stride(0) + 1 * stride(1)\)來計算該元素在陣列中的索引，
除了dtype、device、layout之外，dispatcher還可以用來呼叫legacy operator，比如說addmm這個operator，它的GPU實作就是通過dispatcher來跳轉到legacy::cuda::_th_addmm，
END
到此，就完成了對PyTorch算子的學習，如果你要學習其他算子，可以先從aten/src/ATen/native目錄的相關函式入手，從native_functions.yaml中找到dispatch目標函式，詳情可以參考Figure 1，
更多精彩文章，歡迎掃碼關注下方的公眾號, 并訪問我的簡書博客：https://www.jianshu.com/u/c0fe8671254e
歡迎轉發至朋友圈，作業號轉載請后臺留言申請授權～
轉載請註明出處，本文鏈接：https://www.uj5u.com/qita/10786.html
標籤：其他