譯者序：這篇博文是一篇非常新的介紹PyTorch內(nèi)部機(jī)制的文章，作者Edward Z Yang來自于Stanford大學(xué)，是PyTorch的核心開發(fā)者之一。文章中介紹了如何閱讀PyTorch源碼和擴(kuò)展PyTorch的技巧。目前講PyTorch底層的文章不多，故將其翻譯出來，才疏學(xué)淺，如有疏漏，歡迎留言討論。 原文鏈接：http://blog.ezyang.com/2019/05/pytorch-internals/ 翻譯努力追求通俗、易懂，有些熟知的名詞沒有進(jìn)行翻譯比如(Tensor, 張量) 部分專有名詞翻譯對照表如下 英文 譯文

這篇博文是一篇長論文形式的關(guān)于PyTorch內(nèi)部機(jī)制的演講材料，我于2019年5月14日在PyTorch紐約見面會中進(jìn)行了這場演講。

Intros

大家好！我今天帶來的是關(guān)于PyTorch內(nèi)部機(jī)制的演講

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

這個演講分為兩個部分：在第一部分，我將會向你介紹tensor庫的基本概念。我將會從你所熟知的tensor數(shù)據(jù)類型談起，并且詳細(xì)討論這個數(shù)據(jù)類型提供了什么，作為幫助你理解它是如何實現(xiàn)的指引。如果你是一個PyTorch的重度用戶，大部分內(nèi)容都是你所熟知的。我們也將會討論擴(kuò)展PyTorch的"三要素"：布局(layout)，設(shè)備(device)和數(shù)據(jù)類型(dtype)，這三個要素指導(dǎo)著我們選擇哪種方式擴(kuò)展Tensor類。

在紐約的現(xiàn)場演講中，我跳過了關(guān)于自動微分(autograde)的部分，不過我在這個博文中簡要的討論了它們。 第二部分包含了PyTorch源碼的細(xì)節(jié)。我會告訴你如何在復(fù)雜autograd的代碼中找到邏輯，哪些代碼是重要的，哪些代碼是老舊的，以及所有PyTorch提供的易用工具來幫助你編寫kernels。

Concepts

Tensor/Storage/Strides

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

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

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

假設(shè)我想要訪問位于tensor [1, 0]位置處的元素，如何將這個邏輯地址轉(zhuǎn)化到物理內(nèi)存的地址上呢？步長就是用來解決這樣的問題：當(dāng)我們根據(jù)下標(biāo)索引查找tensor中的任意元素時，將某維度的下標(biāo)索引和對應(yīng)的步長相乘，然后將所有維度乘積相加就可以了。在上圖中我將第一維(行)標(biāo)為藍(lán)色，第二維(列)標(biāo)為紅色，因此你能夠在計算中方便的觀察下標(biāo)和步長的對應(yīng)關(guān)系。

求和返回了一個0維的標(biāo)量2，而內(nèi)存中地址偏移量為2的位置正好儲存了元素3。 (在后面的演講中，我會討論TensorAccessor，一個方便的類來處理下標(biāo)到地址的計算。當(dāng)你使用TensorAccessor而不是原始的指針的時候，這個類能隱藏底層細(xì)節(jié)，自動幫助你完成這樣的計算) 步長是實現(xiàn)PyTorch視圖(view)的根基。例如，假設(shè)我們想要提取上述tensor的第二行：

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

來自于演講的問題：如果我給一個tensor生成了一個視圖，我怎樣釋放掉原tensor的內(nèi)存？ 回答：你必須要復(fù)制一份這個視圖，以切斷和原tensor物理內(nèi)存的關(guān)系。除此之外，別無選擇。順便提一下，如果你之前寫過Java，拿到一個字符串的子字符串有相似的問題，因為默認(rèn)情況下不會產(chǎn)生數(shù)據(jù)的復(fù)制，因此子字符串關(guān)聯(lián)著(可能非常大的)原字符串。這個問題在Java 7u6被修復(fù)了。

一個更有趣的例子是假設(shè)我想要拿第一列的數(shù)據(jù)：

物理內(nèi)存中處于第一列的元素是不連續(xù)的：每個元素之間都隔著一個元素。這里步長就有用武之地了：我們將步長指定為2，表示在當(dāng)前元素和下一個你想訪問的元素之間， 你需要跳躍2個元素(跳過1個元素)。 步長表示法能夠表示所有tensor上有趣的視圖，如果你想要進(jìn)行一些嘗試，見步長可視化。 讓我們退一步想想如何實現(xiàn)這種機(jī)制(畢竟，這是一個關(guān)于內(nèi)部機(jī)制的演講)。要取得tensor上的視圖，我們得對tensor的的邏輯概念和tensor底層的物理數(shù)據(jù)(稱為存儲 storage)進(jìn)行解耦：

一個存儲可能對應(yīng)多個tensor。存儲定義了tensor的數(shù)據(jù)類型和物理大小，而每個tensor記錄了自己的大小(size)，步長(stride)和偏移(offset)，這些元素定義了該tensor如何對存儲進(jìn)行邏輯解釋。 值得注意的是即使對于一些不需要用到存儲的"簡單"的情況(例如，通過torch.zeros(2,2)分配一個內(nèi)存連續(xù)的tensor)，總是存在著Tensor-Storage對。

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

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

第一種分派基于設(shè)備類型(device type)和tensor的布局(layout of a tensor)，例如這個tensor是CPU tensor還是CUDA tensor；或者，這個tensor是基于步長的(strided) tensor 還是稀疏tensor。這是一種動態(tài)分派的過程：使用一個虛函數(shù)調(diào)用實現(xiàn)(虛函數(shù)的細(xì)節(jié)將在教程的后半部分詳述)。這種動態(tài)分派是必要的因為顯然CPU和GPU實現(xiàn)矩陣乘法的方式不同。

這種分派是動態(tài)的因為對應(yīng)的kernels(理解為具體的實現(xiàn)代碼)可能存在于不同的庫中(e.g. libcaffe2.so 或 libcaffe2_gpu.so)，如果你想要訪問一個沒有直接依賴的庫，你就得動態(tài)的分派你的函數(shù)調(diào)用到這些庫中。 第二種分派基于tensor的數(shù)據(jù)類型(dtype)。這種依賴可以通過簡單的switch語句解決。稍稍思考，這種分派也是有必要的：CPU 代碼(或者GPU代碼)實現(xiàn)float類型矩陣乘法和int類型矩陣乘法也會有差異，因此每種數(shù)據(jù)類型(dtype)都需要不同的kernels。 如果你想要理解operators在PyTorch中是如何調(diào)用的，上面這張圖也許最應(yīng)該被記住。當(dāng)講解代碼的時候我們會再回到這張圖。

Layout/Device/Dtype

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

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

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

內(nèi)存布局(The layout) 描述了我們?nèi)绾谓忉屵@些物理內(nèi)存。常見的布局是基于步長的tensor(strided tensor)。稀疏tensor有不同的內(nèi)存布局，通常包含一對tensors，一個用來存儲索引，一個用來存儲數(shù)據(jù)；MKL-DNN tensors 可能有更加不尋常的布局，比如塊布局(blocked layout)，這種布局難以被簡單的步長(strides)表達(dá)。

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

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

還有一種方法可以用來擴(kuò)展Tensor，即寫一個tensor的wrapper類，實現(xiàn)你自己的對象類型(object type)。聽起來很顯然，但是很多人卻在該用wrapper擴(kuò)展的時候卻選擇了擴(kuò)展上述三種要素。wrapper類擴(kuò)展的一個非常好的優(yōu)點是開發(fā)非常簡單。 什么時候我們應(yīng)該寫一個tensor wrapper或者擴(kuò)展PyTorch tensor？一個至關(guān)重要的測試是在反向自動求導(dǎo)的過程中你是否需要傳遞該tensor。

例如通過這樣的測試，我們就可以知道應(yīng)該通過擴(kuò)展PyTorch的方式實現(xiàn)稀疏tensor，而不是建立一個包含索引tensor和值tensor的Python對象(wrapper方式)：因為當(dāng)在一個包含Embedding的網(wǎng)絡(luò)上做優(yōu)化的時候，我們希望生成的梯度也是稀疏的。

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

Autograd

上面講述的都是tensor相關(guān)的東西，不過如果Pytorch僅僅提供了Tensor，那么它不過是numpy的一個克隆。PyTorch 發(fā)布時一個區(qū)別性的特征是提供了自動微分機(jī)制(現(xiàn)在我們有了其他很酷的特性包括TorchScript；但是當(dāng)時，自動微分是僅有的區(qū)別點) 自動微分到底做了什么呢？自動微分是訓(xùn)練神經(jīng)網(wǎng)絡(luò)的一種機(jī)制：

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

請花一點時間學(xué)習(xí)上面這張圖。有一些東西需要展開來講；下面列出了哪些東西值得關(guān)注：

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

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

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

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

在將計算分派到CPU或者CUDA的具體實現(xiàn)之前，變量也要進(jìn)行分派，這個分派的目的是取出變量內(nèi)部封裝的分派函數(shù)的具體實現(xiàn)(上圖中綠色部分)，然后再將結(jié)果封裝到變量里并且為反向計算記錄下必要的自動微分元信息。 當(dāng)然也有其他的實現(xiàn)沒有unwrap操作；他們僅僅調(diào)用其他的變量實現(xiàn)。你可能會花很多時間在變量的調(diào)用棧中跳轉(zhuǎn)。然后，一旦某個變量unwrap并進(jìn)入了非變量的tensor域，變量調(diào)用棧就結(jié)束了，你不會再回到變量域，除非函數(shù)調(diào)用結(jié)束并且返回。

Mechanics

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

PyTorch的源碼包含許多文件目錄，CONTRIBUTING 文件里給這些目錄做了詳細(xì)的解釋，不過實話說，你只需要關(guān)注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 內(nèi)對operators的實現(xiàn)分成兩類，一種是現(xiàn)代的C++實現(xiàn)版本，另一種是老舊的C實現(xiàn)版本，我們不提倡你花太多的時間在C實現(xiàn)的版本上。

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

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

Operator call stack

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

(譯注：下面這一部分需要對C++的機(jī)制有相當(dāng)?shù)牧私猓热缣摵瘮?shù)調(diào)用等等，我添加了一些自己的理解，盡力翻譯得易懂一些，但是不保證完全正確，原文鏈接供參考) 當(dāng)你調(diào)用一個函數(shù)比如torch.add的時候，會發(fā)生哪些事情？如果你記得我們之前討論過的分派機(jī)制，你的腦海中會浮現(xiàn)一個基本的流程：

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

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

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

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

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

上面的C++代碼展示了分派具體怎樣實現(xiàn)的，我們以一個C實現(xiàn)的Python function為例子 (譯注：即下面的THPVariable_add, 以TH開頭的大都是C代碼，后文會介紹)，這種實現(xiàn)在Python代碼中我們會通過類似這樣語句調(diào)用： torch._C.VariableFunctions.add.THPVariable_add。 要強(qiáng)調(diào)的是上面這段代碼是自動生成的。你不會在GitHub repository中搜索到它們，因此你必須得從源碼構(gòu)建PyTorch才能查看到它們。另一個重要的事實是，你不需要深入地了解這段代碼干了什么；簡單的掃一遍代碼并且對大概的思路有個了解就足夠了。

如上圖，我用藍(lán)色標(biāo)注了一些最重要的部分：如你所見，PythonArgParser class 用來從Python (譯注：Python add方法)的 args和kwargs中生成C++ parser對象，(譯注：通過parser對象的parse方法可以得到一個r對象，r里封裝了左操作數(shù)r.tensor(0)，操作符r.scalar(1)和右操作數(shù)r.tensor(1)，見上面的代碼) 然后我們調(diào)用dispatch_add函數(shù)(上圖紅色所示)，它釋放了Python的全局解釋器鎖(global interpreter lock) 然后調(diào)用一個一般方法作用到C++ tensor self上(譯注：self tensor是C++ Tensor類的對象，C++ Tensor類見下面這張圖)。當(dāng)這個方法返回時，我們重新將Tensor封裝回Python object。 (到此為止，ppt上有個疏漏：我應(yīng)該向你展示關(guān)于Variable dispatch的代碼。目前還沒修復(fù)這個部分。你可以想象奇妙的魔法發(fā)生后，我們到了...)

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

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

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

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

Tools for writing kernels

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

我們通常上認(rèn)為一個kernel包含如下部分：

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

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

然后，一般我們會給輸出tensor分配空間，以將結(jié)果寫入進(jìn)去

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

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

最后，你需要訪問數(shù)據(jù)并做希望做的計算！

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

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

你也可能要在 derivatives.yaml 定義operation的求導(dǎo)操作。

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

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

Pytorch中編寫operator的另一件值得注意的事情是，通常對一個operator，你需要編寫三種版本：abs_out這個版本把輸出存儲在(out= 這個關(guān)鍵字參數(shù)中)，abs_這個版本會就地修改輸入，abs這個是常規(guī)版本(返回輸出，輸入不變)。 在大多數(shù)情況下，我們實現(xiàn)的是abs_out版本，然后通過封裝的方式實現(xiàn)abs和abs_，但是也有給每個函數(shù)實現(xiàn)一個單獨版本的時候。

為了做數(shù)據(jù)類型分派(dtype dispatch)，你應(yīng)當(dāng)使用AT_DISPATCH_ALL_TYPES宏。這個宏的輸入?yún)?shù)是Tensor的type，和一個可以分派各種的type類型的lambda表達(dá)式，通常情況下，這個lambda表達(dá)式會調(diào)用一個模板幫助函數(shù)(templated helper function，譯注：也是C++中的概念，C++泛型會討論到模板函數(shù))。 這個宏不僅"做分派工作"，它也決定了你的kernel將會支持哪些數(shù)據(jù)類型。嚴(yán)格來說，這個宏有幾個不同的版本，這些版本可以讓你選擇處理哪些特定的dtype子集。大多數(shù)情況下，你會使用AT_DISPATCH_ALL_TYPES，但是一定要留心當(dāng)你只想要分派到特定類型的場景。關(guān)于在特定場景如何選擇宏詳見Dispatch.h

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

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

如果你僅僅想拿到存儲在特定位置的數(shù)值，你應(yīng)該使用TensorAccessor。tensor accessor類似于tensor，但是它將維度(dimensionality)和數(shù)據(jù)類型(dtype)硬編碼(hard codes)成了模板參數(shù)(template parameters 譯注：代碼里的x.accessor 表示數(shù)據(jù)類型是float, 維度是3)。當(dāng)你通過x.accessor()得到一個accessor實例，PyTorch 會做一個運行時檢測(runtime test)來保證tensor確實是這樣的形式(format，譯注：形式指數(shù)據(jù)類型和維度)，但是在那之后，每一次訪問都不會再檢查。

Tensor accessors能夠正確的處理步長(stride)，因此當(dāng)你做些原始指針(raw pointer)訪問的時候你應(yīng)當(dāng)盡量用它們 (不幸的是，一些老舊的代碼并沒有這樣)。PyTorch里還有一個PackedTensorAccessor類，被用來在CUDA加載過程中傳輸accessor，因此你能夠在CUDA kernel 內(nèi)訪問accessors。(小提示：TensorAccessor默認(rèn)是64-bit索引的，在CUDA中要比32-bit索引要慢很多)

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

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

Legacy code

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

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

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

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

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

Workflow efficiency

作為總結(jié)，我想要討論一些關(guān)于高效擴(kuò)展PyTorch的技巧。如果說龐大的PyTorch C++代碼庫是第一道阻止很多人貢獻(xiàn)代碼到PyTorch的門檻，那么工作效率就是第二道門檻。如果你試著用寫Python的習(xí)慣編寫C++代碼，你將會花費大量的時間，因為重新編譯PyTorch太耗時了，你需要無盡的時間來驗證你的改動是否奏效。 如何高效的改動PyTorch可能需要另一場專門的talk，但是這個PPT總結(jié)了一些常見的"誤區(qū)"：

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

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

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

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

Conclusions

總之這份教程帶你快速掃過PyTorch內(nèi)部機(jī)制！許多東西沒有被討論到，但是希望以上的描述和解釋能夠幫助你對代碼的大體結(jié)構(gòu)有個初步的了解。 看完這份教程后你需要去哪里獲得更詳細(xì)的資源？你能夠做哪種類型的貢獻(xiàn)？一個比較好的起點是我們的問題追蹤器(issue tracker)。在今年早些時候，我們開始對問題進(jìn)行標(biāo)注，一個標(biāo)注過的問題意味著至少有一個PyTorch開發(fā)者注意到了它并且做了初始的任務(wù)評估。

通過這些標(biāo)注你能夠知道我們認(rèn)為哪些問題是high priority的，或者你可以查詢屬于特定模塊的問題，例如 autograd ，或者你可以查詢一些我們認(rèn)為不是那么重要的小問題(警告：我們有時也會判斷失誤) 即使你不想立刻開始編程，也有很多有意義的工作比如改善文檔(我喜歡合并文檔的pull請求，它們實在是太好了)，幫助我們復(fù)現(xiàn)其他用戶報告的bug，幫助我們討論問題追蹤中的RFCs(request for comment，請求給出詳細(xì)注釋)。

責(zé)任編輯：xj

原文標(biāo)題：一文搞懂 PyTorch 內(nèi)部機(jī)制

文章出處：【微信公眾號：深度學(xué)習(xí)自然語言處理】歡迎添加關(guān)注！文章轉(zhuǎn)載請注明出處。