學(xué)習(xí) PyTorch 比較簡(jiǎn)單,，但你能學(xué)習(xí) PyTorch 內(nèi)部機(jī)制嗎？最近，有 14 年 ML 經(jīng)驗(yàn)的大神 Christian 介紹了 PyTorch 的內(nèi)核機(jī)制。雖然在實(shí)際使用中并不需要這些知識(shí)，但探索 PyTorch 內(nèi)核能大大提升我們對(duì)代碼的直覺與理解,，挖底層實(shí)現(xiàn)的都是大神~

PyTorch 的構(gòu)建者表明，Pytorch 的哲學(xué)是解決當(dāng)務(wù)之急,，也就是說即時(shí)構(gòu)建和運(yùn)行我們的計(jì)算圖,。這恰好適合 Python 的編程理念，一邊定義就可以在 Jupyter Notebook 一邊運(yùn)行,，因此,，PyTorch 的工作流程非常接近于 Python 的科學(xué)計(jì)算庫 NumPy。

Christian 表明 PyTorch 之所以這么方便,，很多都是因?yàn)樗摹富颉埂獌?nèi)部運(yùn)行機(jī)制決定的,。這一篇報(bào)告并不會(huì)介紹如何使用 PyTorch 基礎(chǔ)模塊，或如何用 PyTorch 訓(xùn)練一個(gè)神經(jīng)網(wǎng)絡(luò),，Christian 關(guān)注的是如何以直觀的形式介紹 PyTorch 的內(nèi)核機(jī)制,，即各個(gè)模塊到底是怎么工作的。

Christian 在 Reddit 表示這一次報(bào)告由于錄像問題并不能上傳演講視頻,，因此暫時(shí)只能分享演講 PPT,。不過 Christian 最近也會(huì)再做一次該主題的演講，所以我們可以期待下次能有介紹 PyTorch 的視頻,。

演講 PPT 地址：https://speakerdeck.com/perone/pytorch-under-the-hood

百度云地址：https://pan.baidu.com/s/1aaE0I1geF7VwEnQRwmzBtA

如下所示為這次演講的主要議程,，它主要從張量和 JIT 編譯器出發(fā)介紹底層運(yùn)行機(jī)制：

在討論 PyTorch 的各組件機(jī)制前，我們需要了解整體工作流,。PyTorch 使用一種稱之為 imperative / eager 的范式,，即每一行代碼都要求構(gòu)建一個(gè)圖以定義完整計(jì)算圖的一個(gè)部分,。即使完整的計(jì)算圖還沒有完成構(gòu)建,，我們也可以獨(dú)立地執(zhí)行這些作為組件的小計(jì)算圖，這種動(dòng)態(tài)計(jì)算圖被稱為「define-by-run」方法,。

其實(shí)初學(xué)者了解到整體流程就可以學(xué)著使用了,，但底層機(jī)制有助于對(duì)代碼的理解和掌控。

張量

在概念上,，張量就是向量和矩陣的推廣,，PyTorch 中的張量就是元素為同一數(shù)據(jù)類型多維矩陣。雖然 PyTorch 的接口是 Python,，但底層主要都是用 C++實(shí)現(xiàn)的,，而在 Python 中,，集成 C++代碼通常被稱為「擴(kuò)展」。

因?yàn)閺埩恐饕休d數(shù)據(jù),，并進(jìn)行計(jì)算,。PyTorch 的張量計(jì)算使用最底層和基本的張量運(yùn)算庫 ATen，它的自動(dòng)微分使用 Autograd,，該自動(dòng)微分工具同樣建立在 ATen 框架上,。

Python 對(duì)象

為了定義 C/C++中一個(gè)新的 Python 對(duì)象類型，你需要定義如下 THPVariable 類似結(jié)構(gòu),。其中第一個(gè) PyObject_HEAD 宏旨在標(biāo)準(zhǔn)化 Python 對(duì)象,，并擴(kuò)展至另一個(gè)結(jié)構(gòu)，該結(jié)構(gòu)包含一個(gè)指向類型對(duì)象的指針,，以及一個(gè)帶有引用計(jì)數(shù)（ref count）的字段,。

Python API 中有兩個(gè)額外的宏，分別稱為 Py_INCREF() 和 Py_DECREF(),，可用于增加和減少 Python 對(duì)象的引用計(jì)數(shù),。

在 PyThon 中，任何東西都是對(duì)象,，例如變量,、數(shù)據(jù)結(jié)構(gòu)和函數(shù)等。

ZERO-COPYING 張量

由于 Numpy 數(shù)組的使用非常普遍,，我們確實(shí)需要在 Numpy 和 PyTorch 張量之間做轉(zhuǎn)換,。因此 PyTorch 給出了 from_numpy() 和 numpy() 兩個(gè)方法，從而在 NumPy 數(shù)組和 PyTorch 張量之間做轉(zhuǎn)換,。

因?yàn)閺埩績(jī)?chǔ)存的成本比較大,，如果我們?cè)谏鲜鲛D(zhuǎn)換的過程中復(fù)制一遍數(shù)據(jù)，那么內(nèi)存的占用會(huì)非常大,。PyTorch 張量的一個(gè)優(yōu)勢(shì)是它會(huì)保留一個(gè)指向內(nèi)部 NumPy 數(shù)組的指針,，而不是直接復(fù)制它。這意味著 PyTorch 將擁有這一數(shù)據(jù),，并與 NumPy 數(shù)組對(duì)象共享同一內(nèi)存區(qū)域,。

Zero-Copying 的形式確實(shí)能省很多內(nèi)存，但是如上所示在位（in-place）和標(biāo)準(zhǔn)運(yùn)算之間的區(qū)別會(huì)有點(diǎn)模糊,。如果用 np_array = np_array +1.0,，torch_array 的內(nèi)存不會(huì)改變，但是如果用 np_array += 1.0,，torch_array 的內(nèi)存卻又會(huì)改變,。

CPU/GPU 內(nèi)存分配

張量的實(shí)際原始數(shù)據(jù)并不是立即保存在張量結(jié)構(gòu)中，而是保存在我們稱之為「存儲(chǔ)（Storage）」的地方，它是張量結(jié)構(gòu)的一部分,。一般張量存儲(chǔ)可以通過 Allocator 選擇是儲(chǔ)存在計(jì)算機(jī)內(nèi)存（CPU）還是顯存（GPU）,。

THE BIG PICTURE

最后，PyTorch 主張量 THTensor 結(jié)構(gòu)可以展示為下圖,。THTensor 的主要結(jié)構(gòu)為張量數(shù)據(jù),，它保留了 size/strides/dimensions/offsets/等信息，同時(shí)還有存儲(chǔ) THStorage,。

JIT

因?yàn)?PyTorch 是即時(shí)運(yùn)行模式,，這表明它很容易 Debug 或檢查代碼等。在 PyTorch 1.0 中,，其首次引進(jìn)了 torch.jit,，它是一組編譯工具，且主要目標(biāo)是彌補(bǔ)研究與產(chǎn)品部署的差距,。JIT 包含一種名為 Torch Script 的語言,，這種語言是 Python 的子語言。使用 Torch Script 的代碼可以實(shí)現(xiàn)非常大的優(yōu)化,，并且可以序列化以供在后續(xù)的 C++API 中使用,。

如下所示為常見使用 Python 運(yùn)行的 Eager 模式，也可以運(yùn)行 Script 模式,。Eager 模式適合塊做原型與實(shí)驗(yàn),，而 Script 模式適合做優(yōu)化與部署。

那么為什么要用 TORCHSCRIPT 呢,？Christian 給出了以下理由：

PyTorch JIT 主要過程

如下所示 JIT 主要會(huì)輸入代碼或 Python 的抽象句法樹（AST）,，其中 AST 會(huì)用樹結(jié)構(gòu)表征 Python 源代碼的句法結(jié)構(gòu)。解析可能是解析句法結(jié)構(gòu)和計(jì)算圖,，然后語法檢測(cè)接連著代碼優(yōu)化過程,，最后只要編譯并執(zhí)行就可以了。

其中優(yōu)化可以用于模型計(jì)算圖,，例如展開循環(huán)等,。在如下所示的 Peephole 優(yōu)化中，編譯器僅在一個(gè)或多個(gè)基本塊中針對(duì)已生成的代碼,，結(jié)合 CPU 指令的特點(diǎn)和一些轉(zhuǎn)換規(guī)則提升性能,。Peephole 優(yōu)化也可以通過整體分析和指令轉(zhuǎn)換提升代碼性能。

如下所示矩陣的兩次裝置等于矩陣本身,，這應(yīng)該是需要優(yōu)化的,。

執(zhí)行

和 Python 解釋器可以執(zhí)行代碼一樣,，PyTorch 在 JIT 過程中也有一個(gè)解釋器執(zhí)行中間表征指令：

最后,，Christian 還介紹了很多內(nèi)部運(yùn)行機(jī)制，不過因?yàn)樗鼈兌己茈y，而且暫時(shí)沒有提供視頻講解,，讀者大牛們可以看看具體 PPT 內(nèi)容,。