本文以?xún)?nèi)部PyTorch核心開(kāi)發(fā)者視角，非常詳細(xì)透徹的分析了PyTorch的內(nèi)部結(jié)構(gòu)，為開(kāi)發(fā)者提供一張地圖，告訴您“支持自動(dòng)區(qū)分的張量庫(kù)”的基本概念結(jié)構(gòu)，并提供一些工具和技巧，以便找到適合代碼庫(kù)。對(duì)學(xué)習(xí)PyTorch、尤其是致力于參與PyTorch貢獻(xiàn)有非常大的意義。

PyTorch是一個(gè)開(kāi)源的Python機(jī)器學(xué)習(xí)庫(kù)，基于Torch，已成為最受歡迎的機(jī)器學(xué)習(xí)框架之一。

相比Tensorflow，PyTorch的社區(qū)由更多專(zhuān)業(yè)機(jī)器學(xué)習(xí)開(kāi)發(fā)人員、軟件架構(gòu)師和公司內(nèi)部程序員組成。

PyTorch也更多地用于數(shù)據(jù)分析和業(yè)務(wù)環(huán)境中的特殊模型中。

在PyTorch社區(qū)中，有更多的Python開(kāi)發(fā)人員從事Web應(yīng)用程序。此外，這種Python向框架的多功能性，使得研究人員能夠以幾乎無(wú)痛的方式測(cè)試想法，使得它成為最先進(jìn)的尖端解決方案的首選框架。

對(duì)于準(zhǔn)備、正在學(xué)習(xí)PyTorch的讀者來(lái)說(shuō)，了解其內(nèi)部機(jī)制能夠極大的提升學(xué)習(xí)效率、增進(jìn)對(duì)PyTorch設(shè)計(jì)原理和目的的了解，從而能夠更好的在工作學(xué)習(xí)中使用該工具。如果你立志參與到PyTorch后續(xù)的改進(jìn)中，那么更應(yīng)該深入的了解其內(nèi)部機(jī)制。

好消息是，F(xiàn)acebook Research Engineer、斯坦福博士生、PyTorch核心開(kāi)發(fā)人員Edward Z. Yang為大家?guī)?lái)一份PyTorch內(nèi)部機(jī)制的詳解slides，新智元在此強(qiáng)力推薦給廣大讀者。正文約3500字，閱讀可能需要10分鐘。

由于微信的限制無(wú)法展示高清圖像，我們特意為大家在文末找來(lái)了可下載的高清完整版，預(yù)祝大家學(xué)習(xí)愉快！

這份內(nèi)部機(jī)制詳解是為誰(shuí)準(zhǔn)備的？

主要針對(duì)使用過(guò)PyTorch的人，尤其是希望成為PyTorch貢獻(xiàn)者、但卻被PyTorch的龐大復(fù)雜的C++代碼庫(kù)嚇到的人。

最終目的是能夠?yàn)榇蠹姨峁┮粋€(gè)通關(guān)寶典，讓大家了解“支持自動(dòng)區(qū)分的Tensor庫(kù)”的基本概念結(jié)構(gòu)，并提供一些工具和技巧，用來(lái)更容易的找到適合代碼庫(kù)。

讀者只需要對(duì)PyTorch有一個(gè)初步的了解，并且有過(guò)一定的動(dòng)手經(jīng)驗(yàn)即可。門(mén)檻還是非常低的。

全部?jī)?nèi)容分為兩部分。首先介紹Tensor庫(kù)的概念。作者將從Tensor數(shù)據(jù)類(lèi)型開(kāi)始，更詳細(xì)地討論這種數(shù)據(jù)類(lèi)型提供的內(nèi)容，以便讓讀者更好地了解它是如何實(shí)際實(shí)現(xiàn)的。布局、設(shè)備和dtype的三位一體，探討如何考慮對(duì)Tensor類(lèi)的擴(kuò)展。

第二部分將討論P(yáng)yTorch實(shí)戰(zhàn)。例如使用autograd來(lái)降低工作量，哪些代碼關(guān)鍵、為什么？以及各種用來(lái)編寫(xiě)內(nèi)核的超酷的工具。

理解Tensor庫(kù)的概念

Tensor

Tensor是PyTorch中的中心數(shù)據(jù)結(jié)構(gòu)。我們可以將Tensor視為由一些數(shù)據(jù)組成，然后是一些描述Tensor大小的元數(shù)據(jù)，包含元素的類(lèi)型（dtype），Tensor所依賴(lài)的設(shè)備（CPU內(nèi)存？CUDA內(nèi)存？）。以及Strides（步幅）。Strides實(shí)際上是PyTorch的一個(gè)顯著特征。

Tensor是一個(gè)數(shù)學(xué)概念。在計(jì)算機(jī)上最常見(jiàn)的表示是將Tensor中的每個(gè)元素連續(xù)地存儲(chǔ)在內(nèi)存中，將每一行寫(xiě)入內(nèi)存，如上所示。

在上面的例子中，指定Tensor包含32位整數(shù)，每個(gè)整數(shù)位于物理地址中，相互偏移四個(gè)字節(jié)。要記住Tensor的實(shí)際尺寸，還必須記錄哪些尺寸是多余的元數(shù)據(jù)。

假設(shè)我想在邏輯表示中訪問(wèn)位置Tensor[0,1]處的元素。通過(guò)Stride我們應(yīng)該這樣做：

找出Tensor的任何元素存在的位置，將每個(gè)索引乘以該維度的相應(yīng)Stride，并將它們加在一起。

上圖中將第一維藍(lán)色和第二維紅色進(jìn)行了顏色編碼，以便在Stride計(jì)算中跟蹤索引和步幅。

以上是Stride的一個(gè)例子。Stride表示實(shí)際上可以讓你代表Tensor的各種有趣的方法; 如果你想玩弄各種可能性，請(qǐng)查看Stride Visualizer。

可能存在共享相同存儲(chǔ)的多個(gè)Tensor，但請(qǐng)記住一點(diǎn)：有Tensor的地方，就有存儲(chǔ)。

存儲(chǔ)定義Tensor的dtype和物理大小，而每個(gè)Tensor記錄大小，步幅和偏移，定義物理內(nèi)存的邏輯解釋。

Tensor擴(kuò)展

有很多有趣的擴(kuò)展，如XLA張量，量化張量，或MKL-DNN張量，作為張量庫(kù)，我們必須考慮是如何適應(yīng)這些擴(kuò)展。

當(dāng)前的擴(kuò)展模型在張量上提供了四個(gè)擴(kuò)展點(diǎn)。首先，用三個(gè)參數(shù)用來(lái)確定張量是什么：

設(shè)備

張量的物理存儲(chǔ)器實(shí)際存儲(chǔ)在何處，例如在CPU上，NVIDIA GPU（cuda）上，或者可能在AMD GPU（hip）或TPU（xla）上的描述。設(shè)備的顯著特征是它有自己的分配器，不能與任何其他設(shè)備一起使用。

布局

布局用來(lái)描述我們?nèi)绾芜壿嫷亟忉屵@個(gè)物理內(nèi)存。最常見(jiàn)的布局是跨步張量，但稀疏張量具有不同的布局，涉及2個(gè)張量：一個(gè)用于索引、一個(gè)用于數(shù)據(jù)。

MKL-DNN張量可能具有更奇特的布局，例如阻擋布局，這不能僅使用步幅來(lái)表示。

dtype

描述了它實(shí)際存儲(chǔ)在張量的每個(gè)元素中的含義。這可以是浮點(diǎn)數(shù)或整數(shù)，或者它可以是例如量化的整數(shù)。

順便說(shuō)一下，如果你想為PyTorch張量添加一個(gè)擴(kuò)展名，請(qǐng)聯(lián)系PyTorch官方。

實(shí)戰(zhàn)技巧

了解你手里的武器

PyTorch有很多文件夾，CONTRIBUTING文檔有非常詳細(xì)的描述。但實(shí)際上，你真正需要了解的只有四個(gè)：

torch/：包含導(dǎo)入和使用的實(shí)際Python模塊。Python代碼，很容易上手調(diào)試。

torch/csrc/：它實(shí)現(xiàn)了在Python和C++之間進(jìn)行轉(zhuǎn)換的綁定代碼，以及一些非常重要的PyTorch功能，如autograd引擎和JIT編譯器。它還包含C++前臺(tái)代碼。

aten/：“A Tensor Library”的縮寫(xiě)（由Zachary DeVito創(chuàng)造），是一個(gè)實(shí)現(xiàn)Tensors操作的C++庫(kù)。存放一些內(nèi)核代碼存在的地方，盡量不要在那里花太多時(shí)間。

c10/：這是一個(gè)雙關(guān)語(yǔ)。C代表Caffe，10既是二級(jí)制的2，也是十進(jìn)制的10（英文Ten，同時(shí)也是Tensor的前半部分）。包含PyTorch的核心抽象，包括Tensor和Storage數(shù)據(jù)結(jié)構(gòu)的實(shí)際實(shí)現(xiàn)。

讓我們看看這種代碼分離在實(shí)踐中是如何分解的：

調(diào)用一個(gè)函數(shù)的時(shí)候，會(huì)經(jīng)歷以下步驟：

將Python翻譯成C

處理變量調(diào)度

處理設(shè)備類(lèi)型/布局調(diào)度

我們有實(shí)際的內(nèi)核，它既可以是現(xiàn)代本機(jī)函數(shù)，也可以是傳統(tǒng)的TH函數(shù)

值得一提的是，所有代碼都是自動(dòng)生成的，所以不會(huì)出現(xiàn)在GitHub的repo里，必須自己構(gòu)建PyTorch后才能看到。不過(guò)你也不必非常深刻地理解這段代碼在做什么，自動(dòng)生成的嘛。

從武器庫(kù)中挑選寫(xiě)內(nèi)核的趁手兵刃

PyTorch為內(nèi)核編寫(xiě)者提供了許多有用的工具。在本節(jié)中，我們將介紹其中比較趁手的工具。

要利用PyTorch帶來(lái)的所有代碼生成，需要為運(yùn)算符編寫(xiě)schema。詳細(xì)介紹參見(jiàn)GitHub的README。

錯(cuò)誤檢查可以通過(guò)低階API（TORCH_CHECK）和高階API實(shí)現(xiàn)。高階API可以基于TensorArg元數(shù)據(jù)提供用戶(hù)友好的錯(cuò)誤消息。

要執(zhí)行dtype調(diào)度，應(yīng)該使用AT_DISPATCH_ALL_TYPES宏，用來(lái)獲取張量的dtype，并用于可從宏調(diào)度的每個(gè)dtype的lambda。通常，這個(gè)lambda只調(diào)用一個(gè)模板化的輔助函數(shù)。

如何提高工作效率

別編輯header！

編輯header會(huì)導(dǎo)致很長(zhǎng)的重構(gòu)時(shí)間，盡量去編輯.cpp文件。

別直接用CI去測(cè)試

CI是一個(gè)直接可用的測(cè)試代碼的變動(dòng)是否有效的非常棒的工具，但如果你真的一點(diǎn)不都改設(shè)置恐怕要浪費(fèi)很長(zhǎng)時(shí)間在測(cè)試過(guò)程中。

強(qiáng)烈建議設(shè)置ccache

它有可能讓你避免在編輯header時(shí)進(jìn)行大量重新編譯。而當(dāng)我們?cè)诓恍枰匦戮幾g文件時(shí)進(jìn)行了重新編譯，它還有助于掩蓋構(gòu)建系統(tǒng)中的錯(cuò)誤。

用一臺(tái)高性能的工作站

如果你建立一個(gè)帶有CPU和RAM的強(qiáng)大服務(wù)器，你將獲得更愉快的體驗(yàn)。特別是，不建議在筆記本電腦上進(jìn)行CUDA構(gòu)建。