想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

機器之心發表於2019-03-01

學習 PyTorch 比較簡單,但你能學習 PyTorch 內部機制嗎?最近,有 14 年 ML 經驗的大神 Christian 介紹了 PyTorch 的核心機制。雖然在實際使用中並不需要這些知識,但探索 PyTorch 核心能大大提升我們對程式碼的直覺與理解,挖底層實現的都是大神~

PyTorch 的構建者表明,Pytorch 的哲學是解決當務之急,也就是說即時構建和執行我們的計算圖。這恰好適合 Python 的程式設計理念,一邊定義就可以在 Jupyter Notebook 一邊執行,因此,PyTorch 的工作流程非常接近於 Python 的科學計算庫 NumPy

Christian 表明 PyTorch 之所以這麼方便,很多都是因為它的「基因」——內部執行機制決定的。這一篇報告並不會介紹如何使用 PyTorch 基礎模組,或如何用 PyTorch 訓練一個神經網路,Christian 關注的是如何以直觀的形式介紹 PyTorch 的核心機制,即各個模組到底是怎麼工作的。

Christian 在 Reddit 表示這一次報告由於錄影問題並不能上傳演講視訊,因此暫時只能分享演講 PPT。不過 Christian 最近也會再做一次該主題的演講,所以我們可以期待下次能有介紹 PyTorch 的視訊。

如下所示為這次演講的主要議程,它主要從張量和 JIT 編譯器出發介紹底層執行機制:

想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

在討論 PyTorch 的各元件機制前,我們需要了解整體工作流。PyTorch 使用一種稱之為 imperative / eager 的正規化,即每一行程式碼都要求構建一個圖以定義完整計算圖的一個部分。即使完整的計算圖還沒有完成構建,我們也可以獨立地執行這些作為元件的小計算圖,這種動態計算圖被稱為「define-by-run」方法。

想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

其實初學者瞭解到整體流程就可以學著使用了,但底層機制有助於對程式碼的理解和掌控。

張量

在概念上,張量就是向量和矩陣的推廣,PyTorch 中的張量就是元素為同一資料型別多維矩陣。雖然 PyTorch 的介面是 Python,但底層主要都是用 C++實現的,而在 Python 中,整合 C++程式碼通常被稱為「擴充套件」。

因為張量主要承載資料,並進行計算。PyTorch 的張量計算使用最底層和基本的張量運算庫 ATen,它的自動微分使用 Autograd,該自動微分工具同樣建立在 ATen 框架上。

Python 物件

為了定義 C/C++中一個新的 Python 物件型別,你需要定義如下 THPVariable 類似結構。其中第一個 PyObject_HEAD 巨集旨在標準化 Python 物件,並擴充套件至另一個結構,該結構包含一個指向型別物件的指標,以及一個帶有引用計數(ref count)的欄位。

想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

Python API 中有兩個額外的巨集,分別稱為 Py_INCREF() 和 Py_DECREF(),可用於增加和減少 Python 物件的引用計數。

在 PyThon 中,任何東西都是物件,例如變數、資料結構和函式等。

想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

ZERO-COPYING 張量

由於 Numpy 陣列的使用非常普遍,我們確實需要在 Numpy 和 PyTorch 張量之間做轉換。因此 PyTorch 給出了 from_numpy() 和 numpy() 兩個方法,從而在 NumPy 陣列和 PyTorch 張量之間做轉換。

想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

因為張量儲存的成本比較大,如果我們在上述轉換的過程中複製一遍資料,那麼記憶體的佔用會非常大。PyTorch 張量的一個優勢是它會保留一個指向內部 NumPy 陣列的指標,而不是直接複製它。這意味著 PyTorch 將擁有這一資料,並與 NumPy 陣列物件共享同一記憶體區域。

想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

Zero-Copying 的形式確實能省很多記憶體,但是如上所示在位(in-place)和標準運算之間的區別會有點模糊。如果用 np_array = np_array +1.0,torch_array 的記憶體不會改變,但是如果用 np_array += 1.0,torch_array 的記憶體卻又會改變。

CPU/GPU 記憶體分配

張量的實際原始資料並不是立即儲存在張量結構中,而是儲存在我們稱之為「儲存(Storage)」的地方,它是張量結構的一部分。一般張量儲存可以通過 Allocator 選擇是儲存在計算機記憶體(CPU)還是視訊記憶體(GPU)。

想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

THE BIG PICTURE

最後,PyTorch 主張量 THTensor 結構可以展示為下圖。THTensor 的主要結構為張量資料,它保留了 size/strides/dimensions/offsets/等資訊,同時還有儲存 THStorage。

想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

JIT

因為 PyTorch 是即時執行模式,這表明它很容易 Debug 或檢查程式碼等。在 PyTorch 1.0 中,其首次引進了 torch.jit,它是一組編譯工具,且主要目標是彌補研究與產品部署的差距。JIT 包含一種名為 Torch Script 的語言,這種語言是 Python 的子語言。使用 Torch Script 的程式碼可以實現非常大的優化,並且可以序列化以供在後續的 C++API 中使用。

如下所示為常見使用 Python 執行的 Eager 模式,也可以執行 Script 模式。Eager 模式適合塊做原型與實驗,而 Script 模式適合做優化與部署。

想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

那麼為什麼要用 TORCHSCRIPT 呢?Christian 給出了以下理由:

想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

PyTorch JIT 主要過程

如下所示 JIT 主要會輸入程式碼或 Python 的抽象句法樹(AST),其中 AST 會用樹結構表徵 Python 原始碼的句法結構。解析可能是解析句法結構和計算圖,然後語法檢測接連著程式碼優化過程,最後只要編譯並執行就可以了。

想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

其中優化可以用於模型計算圖,例如展開迴圈等。在如下所示的 Peephole 優化中,編譯器僅在一個或多個基本塊中針對已生成的程式碼,結合 CPU 指令的特點和一些轉換規則提升效能。Peephole 優化也可以通過整體分析和指令轉換提升程式碼效能。

如下所示矩陣的兩次裝置等於矩陣本身,這應該是需要優化的。

想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

執行

和 Python 直譯器可以執行程式碼一樣,PyTorch 在 JIT 過程中也有一個直譯器執行中間表徵指令:

想讀讀PyTorch底層程式碼?這份核心機制簡介送給你

最後,Christian 還介紹了很多內部執行機制,不過因為它們都很難,而且暫時沒有提供視訊講解,讀者大牛們可以看看具體 PPT 內容。


相關文章