學習 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 的視頻。

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

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

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

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

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

張量

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

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

Python 對象

為了定義 C/C++中一個新的 Python 對象類型，你需要定義如下 THPVariable 類似結構。其中第一個 PyObject_HEAD 宏旨在標準化 Python 對象，并擴展至另一個結構，該結構包含一個指向類型對象的指針，以及一個帶有引用計數（ref count）的字段。

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

在 PyThon 中，任何東西都是對象，例如變量、數據結構和函數等。

ZERO-COPYING?張量

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

因為張量儲存的成本比較大，如果我們在上述轉換的過程中復制一遍數據，那么內存的占用會非常大。PyTorch?張量的一個優勢是它會保留一個指向內部 NumPy 數組的指針，而不是直接復制它。這意味著 PyTorch 將擁有這一數據，并與 NumPy 數組對象共享同一內存區域。

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

CPU/GPU 內存分配

張量的實際原始數據并不是立即保存在張量結構中，而是保存在我們稱之為「存儲（Storage）」的地方，它是張量結構的一部分。一般張量存儲可以通過 Allocator 選擇是儲存在計算機內存（CPU）還是顯存（GPU）。

THE BIG PICTURE

最后，PyTorch 主張量?THTensor 結構可以展示為下圖。THTensor 的主要結構為張量數據，它保留了 size/strides/dimensions/offsets/等信息，同時還有存儲 THStorage。

JIT

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

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

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

PyTorch JIT 主要過程

如下所示 JIT 主要會輸入代碼或 Python 的抽象句法樹（AST），其中 AST 會用樹結構表征 Python 源代碼的句法結構。解析可能是解析句法結構和計算圖，然后語法檢測接連著代碼優化過程，最后只要編譯并執行就可以了。

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

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

執行

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

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

參考鏈接：https://www.reddit.com/r/MachineLearning/comments/avfoso/p_pytorch_under_the_hood/

總結

以上是生活随笔為你收集整理的想读读PyTorch底层代码？这份内核机制简介送给你的全部內容，希望文章能夠幫你解決所遇到的問題。

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