文 | ArchWalker

---

譯者序：這篇博文是一篇非常新的介紹PyTorch內部機制的文章，作者Edward Z Yang來自于Stanford大學，是PyTorch的核心開發(fā)者之一。文章中介紹了如何閱讀PyTorch源碼和擴展PyTorch的技巧。目前講PyTorch底層的文章不多，故將其翻譯出來，才疏學淺，如有疏漏，歡迎留言討論。

原文鏈接：
http://blog.ezyang.com/2019/05/pytorch-internals/

Intros

大家好！我今天帶來的是關于PyTorch內部機制的演講

這個演講的受眾是那些已經用過PyTorch，問過自己"如果我能給PyTorch做些貢獻豈不美哉"但是又被PyTorch龐大的C++代碼嚇退的人。實話說：有時候PyTorch的代碼庫確實又多又雜。這個演講的目的是給你提供一張導向圖：告訴你PyTorch這個"支持自動微分的tensor庫"的基本結構，給你介紹一些能夠幫助你在PyTorch代碼庫中暢游的工具和小技巧。我假設你之前寫過一些PyTorch代碼，但是不需要你對如何實現(xiàn)一個機器學習庫有過深入的理解。

這個演講分為兩個部分：在第一部分，我將會向你介紹tensor庫的基本概念。我將會從你所熟知的tensor數據類型談起，并且詳細討論這個數據類型提供了什么，作為幫助你理解它是如何實現(xiàn)的指引。如果你是一個PyTorch的重度用戶，大部分內容都是你所熟知的。我們也將會討論擴展PyTorch的"三要素"：布局(layout)，設備(device)和數據類型(dtype)，這三個要素指導著我們選擇哪種方式擴展Tensor類。在紐約的現(xiàn)場演講中，我跳過了關于自動微分(autograde)的部分，不過我在這個博文中簡要的討論了它們。

第二部分包含了PyTorch源碼的細節(jié)。我會告訴你如何在復雜autograd的代碼中找到邏輯，哪些代碼是重要的，哪些代碼是老舊的，以及所有PyTorch提供的易用工具來幫助你編寫kernels。

Concepts

Tensor/Storage/Strides

Tensor 是PyTorch的核心數據結構。你可能對tensor的概念已經相當了解了：它是包含若干個標量(標量可以是各種數據類型如浮點型、整形等)的n-維的數據結構。我們可以認為tensor包含了數據和元數據(metadata)，元數據用來描述tensor的大小、其包含內部數據的類型、存儲的位置(CPU內存或是CUDA顯存?)

也有一些你可能不太熟悉的元數據：步長(the stride)，步長實際上是PyTorch的一個亮點，所以值得花點時間好好討論一下它。

Tensor是一個數學概念。當用計算機表示數學概念的時候，通常我們需要定義一種物理存儲方式。最常見的表示方式是將Tensor中的每個元素按照次序連續(xù)的在內存中鋪開(這是術語contiguous的來歷)，將每一行寫到相應內存位置里。如上圖所示，假設tensor包含的是32位的整數，因此每個整數占據一塊物理內存，每個整數的地址都和上下相鄰整數相差4個字節(jié)。為了記住tensor的實際維度，我們需要將tensor的維度大小記錄在額外的元數據中。

那么，步長在物理表示中的作用是什么呢？

假設我想要訪問位于tensor [1, 0]位置處的元素，如何將這個邏輯地址轉化到物理內存的地址上呢？步長就是用來解決這樣的問題：當我們根據下標索引查找tensor中的任意元素時，將某維度的下標索引和對應的步長相乘，然后將所有維度乘積相加就可以了。在上圖中我將第一維(行)標為藍色，第二維(列)標為紅色，因此你能夠在計算中方便的觀察下標和步長的對應關系。求和返回了一個0維的標量2，而內存中地址偏移量為2的位置正好儲存了元素3。

(在后面的演講中，我會討論TensorAccessor，一個方便的類來處理下標到地址的計算。當你使用TensorAccessor而不是原始的指針的時候，這個類能隱藏底層細節(jié)，自動幫助你完成這樣的計算)

步長是實現(xiàn)PyTorch視圖(view)的根基。例如，假設我們想要提取上述tensor的第二行：

使用高級索引技巧，我只需要寫成tensor[1, :] 來獲取這一行。重要的事情是：這樣做沒有創(chuàng)建一個新的tensor；相反，它只返回了原tensor底層數據的另一個視圖。這意味著如果我編輯了這個視圖中的數據，變化也會反應到原tensor上。在這個例子中，不難看出視圖是怎么做的：3和4存儲在連續(xù)的內存中，我們所要做的是記錄一個偏移量(offset)，用來表示新的視圖的數據開始于原tensor數據自頂向下的第二個。(每一個tensor都會記錄一個偏移量，但是大多數時候他們都是0，我在圖片中忽略了這樣的例子)

來自于演講的問題：如果我給一個tensor生成了一個視圖，我怎樣釋放掉原tensor的內存？

回答：你必須要復制一份這個視圖，以切斷和原tensor物理內存的關系。除此之外，別無選擇。順便提一下，如果你之前寫過Java，拿到一個字符串的子字符串有相似的問題，因為默認情況下不會產生數據的復制，因此子字符串關聯(lián)著(可能非常大的)原字符串。這個問題在Java 7u6被修復了。

一個更有趣的例子是假設我想要拿第一列的數據：

物理內存中處于第一列的元素是不連續(xù)的：每個元素之間都隔著一個元素。這里步長就有用武之地了：我們將步長指定為2，表示在當前元素和下一個你想訪問的元素之間， 你需要跳躍2個元素(跳過1個元素)。

步長表示法能夠表示所有tensor上有趣的視圖，如果你想要進行一些嘗試，見步長可視化。

讓我們退一步想想如何實現(xiàn)這種機制(畢竟，這是一個關于內部機制的演講)。要取得tensor上的視圖，我們得對tensor的的邏輯概念和tensor底層的物理數據(稱為存儲 storage)進行解耦：

一個存儲可能對應多個tensor。存儲定義了tensor的數據類型和物理大小，而每個tensor記錄了自己的大小(size)，步長(stride)和偏移(offset)，這些元素定義了該tensor如何對存儲進行邏輯解釋。

值得注意的是即使對于一些不需要用到存儲的"簡單"的情況(例如，通過torch.zeros(2,2)分配一個內存連續(xù)的tensor)，總是存在著Tensor-Storage對。

順便提一下，我們也對改進這樣的模型很感興趣。相比于有一個獨立的存儲，只基于現(xiàn)有tensor定義一個視圖。這有一點點復雜，但是優(yōu)點是可以更加直接的表示連續(xù)tensor，而不需要tensor到存儲的轉化。這樣的變化將會使PyTorch的內部表示更加像Numpy。

我們對于tensor的數據布局(data layout)做了相當多的討論，(有人會說，如果你能夠將數據底層表示搞清楚，剩下的一切就順理成章了）。但是我覺得還是有必要簡要的探討一下tensor上的操作(operations)是如何實現(xiàn)的。抽象來說，當你調用torch.mm的時候，會產生兩種分派(dispatch)：

第一種分派基于設備類型(device type)和tensor的布局(layout of a tensor)，例如這個tensor是CPU tensor還是CUDA tensor；或者，這個tensor是基于步長的(strided) tensor 還是稀疏tensor。這是一種動態(tài)分派的過程：使用一個虛函數調用實現(xiàn)(虛函數的細節(jié)將在教程的后半部分詳述)。這種動態(tài)分派是必要的因為顯然CPU和GPU實現(xiàn)矩陣乘法的方式不同。這種分派是動態(tài)的因為對應的kernels(理解為具體的實現(xiàn)代碼)可能存在于不同的庫中(e.g. libcaffe2.so 或 libcaffe2_gpu.so)，如果你想要訪問一個沒有直接依賴的庫，你就得動態(tài)的分派你的函數調用到這些庫中。

第二種分派基于tensor的數據類型(dtype)。這種依賴可以通過簡單的switch語句解決。稍稍思考，這種分派也是有必要的：CPU 代碼(或者GPU代碼)實現(xiàn)float類型矩陣乘法和int類型矩陣乘法也會有差異，因此每種數據類型(dtype)都需要不同的kernels。

如果你想要理解operators在PyTorch中是如何調用的，上面這張圖也許最應該被記住。當講解代碼的時候我們會再回到這張圖。

Layout/Device/Dtype

既然我們一直在討論Tensor，我還想花點時間討論下tensor擴展(extension)。畢竟，日常生活中遇到的tensor大部分都并不是稠密的浮點數tensor。很多有趣的擴展包括XLA tensors，quantized tensors，或者MKL-DNN tensors。作為一個tensor library我們需要考慮如何融合各種類型的tensors。

目前來說PyTorch的擴展模型提供了4種擴展方法。首先，能夠唯一確定Tensor類型的"三要素"是：

• 設備類型(The device) 設備類型描述了tensor的到底存儲在哪里，比如在CPU內存上還是在NVIDIA GPU顯存上，在AMD GPU(hip)上還是在TPU(xla)上。不同設備的特征是它們有自己的存儲分配器(allocator)，不同設備的分配器不能混用。

• 內存布局(The layout) 描述了我們如何解釋這些物理內存。常見的布局是基于步長的tensor(strided tensor)。稀疏tensor有不同的內存布局，通常包含一對tensors，一個用來存儲索引，一個用來存儲數據；MKL-DNN tensors 可能有更加不尋常的布局，比如塊布局(blocked layout)，這種布局難以被簡單的步長(strides)表達。

• 數據類型(The dtype) 數據類型描述tensor中的每個元素如何被存儲的，他們可能是浮點型或者整形，或者量子整形。

如何你想要增加一種PyTorch tensor類型(順便說下，請聯(lián)系我們如果你真的想要做這個！這個目前來說不是那么容易的事情)，你應該想想你要擴展上面提到的哪一個決定張量類型的因素("三要素")。目前為止，并不是所有的組合都有對應的kernel(比如FPGA上稀疏量子張量的計算就沒有現(xiàn)成的kernel)，但是原則上來說大部分的組合都可能是道理的，因此至少在一定程度上我們支持它們。

還有一種方法可以用來擴展Tensor，即寫一個tensor的wrapper類，實現(xiàn)你自己的對象類型(object type)。聽起來很顯然，但是很多人卻在該用wrapper擴展的時候卻選擇了擴展上述三種要素。wrapper類擴展的一個非常好的優(yōu)點是開發(fā)非常簡單。

什么時候我們應該寫一個tensor wrapper或者擴展PyTorch tensor？一個至關重要的測試是在反向自動求導的過程中你是否需要傳遞該tensor。例如通過這樣的測試，我們就可以知道應該通過擴展PyTorch的方式實現(xiàn)稀疏tensor，而不是建立一個包含索引tensor和值tensor的Python對象(wrapper方式)：因為當在一個包含Embedding的網絡上做優(yōu)化的時候，我們希望生成的梯度也是稀疏的。

我們關于tensor擴展的哲學也對tensor自身的數據布局產生著一定的影響。我們始終希望tensor結構能有個固定的布局：我們不希望一些基礎的operator(這些operator經常被調用)，如size of tensor需要一個虛分派 (virtual dispatches)。因此當你觀察Tensor實際的布局的時候(定義在 TensorImpl 結構體中)，一些被我們認為是所有類型tensor都會有的字段定義在前面，隨后跟著一些strided tensors特有的字段(我們也認為它們很重要)，最后才是特定類型tensor的獨有字段，比如稀疏tensor的索引和值。

Autograd

上面講述的都是tensor相關的東西，不過如果Pytorch僅僅提供了Tensor，那么它不過是numpy的一個克隆。PyTorch 發(fā)布時一個區(qū)別性的特征是提供了自動微分機制(現(xiàn)在我們有了其他很酷的特性包括TorchScript；但是當時，自動微分是僅有的區(qū)別點)

自動微分到底做了什么呢？自動微分是訓練神經網絡的一種機制：

…下面這張圖補充了計算loss的gradients所需要的代碼：

請花一點時間學習上面這張圖。有一些東西需要展開來講；下面列出了哪些東西值得關注：

首先請忽略掉那些紅色和藍色的代碼。PyTorch實現(xiàn)了reverse-mode automatic differentiation (反向模式自動微分)，意味著我們通過反向遍歷計算圖的方式計算出梯度。注意看變量名：我們在紅色代碼區(qū)域的最下面計算了loss；然后，在藍色代碼區(qū)域首先我們計算了grad_loss。loss 由 next_h2計算而來，因此我們計算grad_next_h2。嚴格來講，這些以grad_開頭的變量其實并不是gradients；他們實際上是Jacobian矩陣左乘了一個向量，但是在PyTorch中我們就叫它們grad，大部分人都能理解其中的差異。

即使代碼結構相同，代碼的行為也是不同的：前向(forwards)的每一行被一個微分計算代替，表示對這個前向操作的求導。例如，tanh操作符變成了tanh_backward操作符(如上圖最左邊的綠線所關聯(lián)的兩行所示)。前向和后向計算的輸入和輸出顛倒過來：如果前向操作生成了next_h2，那么后向操作取grad_next_h2作為輸入。

概述之，自動微分做了下圖所示的計算，不過實質上沒有生成執(zhí)行這些計算所需的代碼。PyTorch 自動微分不會做代碼到代碼的轉換工作(即使PyTorch JIT確實知道如何做符號微分(symbolic differentiation))。

為了實現(xiàn)這個，當我們在tensor上調用各種operations的時候，一些元數據(metadata)也需要被記錄下來。讓我們調整一下tensor數據結構的示意圖：現(xiàn)在不僅僅單單一個tensor指向storage，我們會有一個封裝著這個tensor和更多信息(自動微分元信息(AutogradeMeta))的變量(variable)。這個變量所包含的信息是用戶調用loss.backward()執(zhí)行自動微分所必備的。

順便我們也更新下分派的圖：

在將計算分派到CPU或者CUDA的具體實現(xiàn)之前，變量也要進行分派，這個分派的目的是取出變量內部封裝的分派函數的具體實現(xiàn)(上圖中綠色部分)，然后再將結果封裝到變量里并且為反向計算記錄下必要的自動微分元信息。

當然也有其他的實現(xiàn)沒有unwrap操作；他們僅僅調用其他的變量實現(xiàn)。你可能會花很多時間在變量的調用棧中跳轉。然后，一旦某個變量unwrap并進入了非變量的tensor域，變量調用棧就結束了，你不會再回到變量域，除非函數調用結束并且返回。

Mechanics

到此我們已經討論了足夠的概念了，現(xiàn)在來看看具體的代碼實現(xiàn)。

PyTorch的源碼包含許多文件目錄，CONTRIBUTING 文件里給這些目錄做了詳細的解釋，不過實話說，你只需要關注4個目錄：

• 首先，torch/包含了你最熟悉的部分：你在代碼中引入并使用的Python 模塊(modules)，這里都是Python代碼，容易修改起來做各種小實驗，然后，暗藏在這些表層代碼的下面是：

• torch/csrc/，這部分C++代碼實現(xiàn)了所謂的PyTorch前端(the frontend of PyTorch)。具體來說，這一部分主要橋接了Python邏輯的C++的實現(xiàn)，和一些PyTorch中非常重要的部分，比如自動微分引擎(autograd engine)和JIT編譯器(JIT compiler)。

• aten/，是"A Tensor Library"的縮寫，是一個C++庫實現(xiàn)了Tensor的各種operations。如果你需要查找一些實現(xiàn)kernels的代碼，很大幾率上他們在aten/文件夾里。ATen 內對operators的實現(xiàn)分成兩類，一種是現(xiàn)代的C++實現(xiàn)版本，另一種是老舊的C實現(xiàn)版本，我們不提倡你花太多的時間在C實現(xiàn)的版本上。

• c10/ ，是一個來自于Caffe2 和 A”Ten“的雙關語(Caffe 10)，其中包含了PyTorch的核心抽象，Tensor和Storage數據結構的實際實現(xiàn)部分。

有如此多的地方看源碼，我們也許應該精簡一下目錄結構，但目前就是這樣。如果你做一些和operators相關的工作，你將花大部分時間在aten上。

Operator call stack

下面我們來看看實踐中這些分離的代碼分別用在那些地方：

(譯注：下面這一部分需要對C++的機制有相當的了解，比如虛函數調用等等，我添加了一些自己的理解，盡力翻譯得易懂一些，但是不保證完全正確，原文鏈接供參考)

當你調用一個函數比如torch.add的時候，會發(fā)生哪些事情？如果你記得我們之前討論過的分派機制，你的腦海中會浮現(xiàn)一個基本的流程：

我們將會從Python 代碼轉到 C++代碼(通過解析Python調用的參數) (譯注：解析調用參數下面代碼中有例子)

處理變量分派(VariableType到Type)，順便說一下，這里的Type和程序語言類型沒有關系，只是在分派中我們這么叫它) (譯注：這一部分博文中沒有討論，下面作者也澄清了這是個疏忽，所以忽略就好了)

處理 設備類型/布局 分派(Type) (譯注：這一部分討論)

找到實際上的kernel，可能是一個現(xiàn)代的函數(modern native funciton)，可能是一個老舊的函數(legacy TH funciton, TH 后面會解釋) (譯注：現(xiàn)代的函數指C++代碼，老舊的多指C代碼，后面有詳細討論。)

每一個步驟具體對應到一些代碼。讓我們剖析這一部分代碼：

上面的C++代碼展示了分派具體怎樣實現(xiàn)的，我們以一個C實現(xiàn)的Python function為例子 (譯注：即下面的THPVariable_add, 以TH開頭的大都是C代碼，后文會介紹)，這種實現(xiàn)在Python代碼中我們會通過類似這樣語句調用：

torch._C.VariableFunctions.add.THPVariable_add。

要強調的是上面這段代碼是自動生成的。你不會在GitHub repository中搜索到它們，因此你必須得從源碼構建PyTorch才能查看到它們。另一個重要的事實是，你不需要深入地了解這段代碼干了什么；簡單的掃一遍代碼并且對大概的思路有個了解就足夠了。如上圖，我用藍色標注了一些最重要的部分：如你所見，PythonArgParser class 用來從Python (譯注：Python add方法)的 args和kwargs中生成C++ parser對象，(譯注：通過parser對象的parse方法可以得到一個r對象，r里封裝了左操作數r.tensor(0)，操作符r.scalar(1)和右操作數r.tensor(1)，見上面的代碼) 然后我們調用dispatch_add函數(上圖紅色所示)，它釋放了Python的全局解釋器鎖(global interpreter lock) 然后調用一個一般方法作用到C++ tensor self上(譯注：self tensor是C++ Tensor類的對象，C++ Tensor類見下面這張圖)。當這個方法返回時，我們重新將Tensor封裝回Python object。

(到此為止，ppt上有個疏漏：我應該向你展示關于Variable dispatch的代碼。目前還沒修復這個部分。你可以想象奇妙的魔法發(fā)生后，我們到了...)

當我們調用C++ Tensor類的add方法時候，虛分派還未發(fā)生。然而，一個內聯(lián)(inline)函數會在"Type"對象上調用一個虛函數(譯注：Type對象指代碼中的type()返回的對象，虛函數指add方法)。這個方法才是真正的虛函數(這就是為什么我之前說Type是一個媒介，作用是引出虛調用)。在這個例子里，這個虛函數調用被分派到TypeDefault的類的add實現(xiàn)上，原因是我們提供了一個add的實現(xiàn)，這種實現(xiàn)在任何一種設備類型上(包括CPU和CUDA)都一致(譯注：所以叫TypeDefault)；假如我們對不同的設備有具體的實現(xiàn)，可能會調用類似于CPUFloatType::add這樣的函數，意味著虛函數add最后將實際的add操作分派到的CPU上浮點數相加的具體kernel代碼上。

根據預期，這個PPT很快將會過時了，Roy Li正在做一些替代Type分派的工作，這些工作將會使PyTorch對于移動設備支持的更好。

值得一提的是，所有的代碼，直到對于具體kernel的調用，都是自動生成的。

這里有點繞，所以一旦你對執(zhí)行流程的大方向有一定的了解，我建議你直接跳到kernels的部分。

Tools for writing kernels

PyTorch為kernels編寫者提供了許多實用的工具。在這一節(jié)里，我們將會簡要了解他們之中的一部分。但是首先，一個kernel包含哪些東西？

我們通常上認為一個kernel包含如下部分：

首先，我們?yōu)閗ernel寫了一些元數據(metadata)，這些元數據驅動了代碼生成，讓你不用寫一行代碼就可以在Python中調用kernel。

一旦你訪問了kernel，意味著你經過了設備類型/布局類型的虛函數分派流程。首先你要寫的一點是錯誤檢測(error checking)，以保證輸入tensors有正確的維度。(錯誤檢測非常重要！千萬別跳過它們！)

然后，一般我們會給輸出tensor分配空間，以將結果寫入進去

接下來是編寫合適的kernel。到這里，你應該做數據類型分派(第二種分派類型dtype)，以跳轉到一個為特定數據類型編寫的kernel上。(通常你不用太早做這個，因為可能會產生一些重復的代碼，比如說一些邏輯在任何case上都適用)

許多高效的kernel需要一定程度上的并行，因此你需要利用多核(multi-CPU)系統(tǒng)。(CUDA kernels 暗含著并行的邏輯，因為它的編程模型是建立在大量的并行體系上的)

最后，你需要訪問數據并做希望做的計算！

在接下來的PPT里，我會帶你了解PyTorch提供的一些工具幫助你實現(xiàn)上述步驟。

為了充分利用PyTorch帶來的代碼生成機制，你需要為operator寫一個schema。這個schema需要給定你定義函數的簽名(signature)，并且控制是否我們生成Tensor方法(比如 t.add())以及命名空間函數(比如at::add())。你也需要在schema中指明當一個設備/布局的組合給定的時候，operator的哪一種實現(xiàn)需要被調用。具體格式細節(jié)查看README in native

你也可能要在 derivatives.yaml 定義operation的求導操作。

錯誤檢測既能通過底層API也能通過高層API來實現(xiàn)。底層API如宏(macro)：TORCH_CHECK，輸入一個boolean表達式，跟著一個字符串，如果根據Boolean表達式判斷結果為false，這個宏就會輸出字符串。這個宏比較好的地方是你能將字符串和非字符串數據混合起來輸出，所有的變量都通過他們實現(xiàn)的<<操作符格式化，PyTorch中大多數重要的數據類型都預定義了<<操作符。(譯注：這是C++中字符格式輸出的方式，即通過重載<<操作符)

高層API能夠幫你避免寫重復的錯誤提示。它的工作方式是首先你將每個Tensor封裝進TensorArg中，TensorArg包含這個Tensor的來源信息(比如，通過它的參數名)。然后它提供了一系列封裝好的函數來做各種屬性的檢測；比如，checkDim()用來檢測是否tensor的維度是一個固定的數。如果它不是，這個函數會基于TensorArg中的元數據提供一個可讀性好的錯誤提示。

Pytorch中編寫operator的另一件值得注意的事情是，通常對一個operator，你需要編寫三種版本：abs_out這個版本把輸出存儲在(out= 這個關鍵字參數中)，abs_這個版本會就地修改輸入，abs這個是常規(guī)版本(返回輸出，輸入不變)。

在大多數情況下，我們實現(xiàn)的是abs_out版本，然后通過封裝的方式實現(xiàn)abs和abs_，但是也有給每個函數實現(xiàn)一個單獨版本的時候。

為了做數據類型分派(dtype dispatch)，你應當使用AT_DISPATCH_ALL_TYPES宏。這個宏的輸入參數是Tensor的type，和一個可以分派各種的type類型的lambda表達式，通常情況下，這個lambda表達式會調用一個模板幫助函數(templated helper function，譯注：也是C++中的概念，C++泛型會討論到模板函數)。

這個宏不僅"做分派工作"，它也決定了你的kernel將會支持哪些數據類型。嚴格來說，這個宏有幾個不同的版本，這些版本可以讓你選擇處理哪些特定的dtype子集。大多數情況下，你會使用AT_DISPATCH_ALL_TYPES，但是一定要留心當你只想要分派到特定類型的場景。關于在特定場景如何選擇宏詳見Dispatch.h

在CPU上， 你經常想要并行化你的代碼。在之前，OpenMP 原語(pragmas) 經常被用來做并行化的工作。

在我們需要訪問數據的時候，PyTorch提供了不少選擇。

如果你僅僅想拿到存儲在特定位置的數值，你應該使用TensorAccessor。tensor accessor類似于tensor，但是它將維度(dimensionality)和數據類型(dtype)硬編碼(hard codes)成了模板參數(template parameters 譯注：代碼里的x.accessor<float, 3> 表示數據類型是float, 維度是3)。當你通過x.accessor<float, 3>()得到一個accessor實例，PyTorch 會做一個運行時檢測(runtime test)來保證tensor確實是這樣的形式(format，譯注：形式指數據類型和維度)，但是在那之后，每一次訪問都不會再檢查。Tensor accessors能夠正確的處理步長(stride)，因此當你做些原始指針(raw pointer)訪問的時候你應當盡量用它們 (不幸的是，一些老舊的代碼并沒有這樣)。PyTorch里還有一個PackedTensorAccessor類，被用來在CUDA加載過程中傳輸accessor，因此你能夠在CUDA kernel 內訪問accessors。(小提示：TensorAccessor默認是64-bit索引的，在CUDA中要比32-bit索引要慢很多)

如果你編寫的operator需要做一些規(guī)律性的數據訪問，比如，點乘操作，強烈建議你用高層API比如TensorIterator。這個幫助類自動幫你處理了廣播(broadcasting)和類型提升(type promotion)，非常方便。(譯注：廣播和類型提升可以參考numpy相關的描述)

為了在CPU上執(zhí)行得盡量快，也許你需要使用向量化的CPU指令(vectorized CPU instructions)來編寫kernel。我們也提供了工具！Vec256 類表示一個向量，并提供了一系列的方法以對其向量化的操作(vectorized operations)。幫助函數比如binary_kernel_vec 讓你更加容易得運行向量化的操作，以處理原始的CPU指令不容易處理的向量化的場景。同時，這個類還負責針對不同的指令集編譯不同的kernel，然后在運行時對你CPU所支持的指令集做測試，以使用最合適的kernel。

Legacy code

PyTorch 中的許多kernel仍然由古老的TH類型的代碼實現(xiàn)(順便說一下，TH代表TorcH。縮寫固然很好，但是太常見了，如果你看到了TH，就把它當做老舊的就好了)。下面詳細解釋下什么是老舊的TH類型：

它由C代碼編寫，沒有(或者極少)用到C++

它是由手動引用計數的(當不再使用某個tensor的時候，通過手工調用THTensor_free方法來減少引用計數)

它存在于 generic/文件夾中，意味著我們需要通過定義不同的#define scalar_t來多次編譯。

這些代碼是很"瘋狂"的，我們也不愿意維護它們，所以請不要再向里面添加東西了。你可以做的更有意義的事情是，如果你喜歡編程但是不熟悉關于kernel的編寫，你可以嘗試著移植這些TH函數到ATen里面去。

Workflow efficiency

作為總結，我想要討論一些關于高效擴展PyTorch的技巧。如果說龐大的PyTorch C++代碼庫是第一道阻止很多人貢獻代碼到PyTorch的門檻，那么工作效率就是第二道門檻。如果你試著用寫Python的習慣編寫C++代碼，你將會花費大量的時間，因為重新編譯PyTorch太耗時了，你需要無盡的時間來驗證你的改動是否奏效。

如何高效的改動PyTorch可能需要另一場專門的talk，但是這個PPT總結了一些常見的"誤區(qū)"：

如果你編輯了一個頭文件，尤其是那種包含許多源文件(尤其是包含了CUDA文件)，那么你可能會需要一個非常長時間的重新編譯。為了避免這個，盡量保持只修改cpp文件，盡量少修改頭文件！

我們的CI(譯注：應該指一個云端的已配置好的環(huán)境，見鏈接)是一個非常好的，不需要任何配置的環(huán)境來測試你的修改是否會奏效。但是在你得到結果之前估計需要1到2小時。如果你的修改需要大量的實驗驗證，把時間花在設置一個本地開發(fā)環(huán)境上吧。同樣，如果你遇到了一個特別難以debug的問題，在本地環(huán)境中測試它。你可以下載并且使用我們的Docker鏡像 download and run the Docker images locally

如何貢獻的文檔詳述了如何設置ccache，我們強烈推薦這個，因為很多情況下它會在你修改頭文件時幫助你節(jié)省重新編譯的時間。它也能幫助你避免一些我們編譯系統(tǒng)的bugs，比如重新編譯了一些不該重新編譯的文件。

我們有大量的C++代碼，推薦你在一個有著充足CPU和RAM資源的服務器上編譯。強烈不建議你用自己的筆記本編譯CUDA，編譯CUDA是特特特特別慢的，筆記本不具備快速編譯的能力。

Conclusions

總之這份教程帶你快速掃過PyTorch內部機制！許多東西沒有被討論到，但是希望以上的描述和解釋能夠幫助你對代碼的大體結構有個初步的了解。

后臺回復關鍵詞【入群】

加入賣萌屋NLP/IR/Rec與求職討論群

后臺回復關鍵詞【頂會】

獲取ACL、CIKM等各大頂會論文集！

總結

以上是生活随笔為你收集整理的一文搞懂 PyTorch 内部机制的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。