NeurIPS頂會接收,PyTorch官方論文首次曝光完整設計思路

機器之心發表於2019-12-02

目前,這篇論文已經被 NeurlPS 2019 接收為 poster 論文。

論文地址:https://papers.nips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf

NeurIPS頂會接收,PyTorch官方論文首次曝光完整設計思路

NeurIPS頂會接收,PyTorch官方論文首次曝光完整設計思路

論文放出後,吃瓜群眾紛紛流下了感動的淚水,表示:「以後終於有PyTorh的論文可以引用了」,「我們這些年遺漏的參考文獻(終於放出來了)。」

論文概覽

近年來,隨著深度學習的興起,機器學習工具也如雨後春筍般發展起來。Caffe、CNTK、TensorFlow、Theano 等很多流行框架都構建了一個表徵計算的靜態資料流圖,這些圖可以重複應用於批次資料。
但一直以來,這些深度學習框架要麼關注易用性,要麼關注速度,很難二者兼顧。Pytorch 的出現打破了這一限制。作為一個機器學習庫,它提供了一種命令式、Python 式的程式設計風格,支援程式碼作為模型,使得除錯變得簡單,並且與其他流行的科學計算庫保持一致,同時保持高效並支援 GPU 等硬體加速器。

在這篇論文中,作者介紹了驅動 PyTorch 實現的詳細原則以及這些原則在 PyTorch 架構中的反映。此外,作者還解釋瞭如何謹慎而務實地實現 PyTorch 執行時的關鍵元件,使得這些元件能夠協調配合,達到令人滿意的效能。研究者在幾個常見的基準上展示了 PyTorch單個子系統的效率以及整體速度。

PyTorch 的誕生背景

對於深度學習來說,科學計算領域的四大趨勢正變得越來越重要。

  1. 首先,1960年代,APL、MATLAB、R 和 Julia 等領域特定語言的發展將多維陣列(張量)轉換為由一組複雜數學運算元支援的一級目標,以此對其進行處理。另外,NumPy、Torch、Eigen、Lush 等庫的出現使得基於陣列的程式設計在 Python、Lisp、C++、Lua 等通用語言中變得更加高效。

  2. 其次,自動微分的發展使得導數的繁冗計算可以完全自動化。autograd 包的出現推動了這一技術在 NumPy 陣列中的使用,類似的方法也應用於Chainer、DyNet、Lush、Torch、Jax、Flux.jl 等框架。

  3. 再次,隨著開源軟體運動的發展,科學界已經從 Matlab 等封閉私有軟體轉向開源 Python 生態,後者包含 NumPy、SciPy、Pandas 等庫。

  4. 最後,GPU 等通用大規模並行硬體的出現和商業化提供了深度學習方法所需的算力。

PyTorch 迎合了這些趨勢,它提供了一個由 GPU 加速的、基於陣列的程式設計模型,並透過整合在 Python 生態系統中的自動微分實現可微分。

以可用性為中心的設計

PyTorch的設計理念相對較新,從易用性、可擴充套件性的角度進行了設計。


深度學習模型都只是 Python 程式

神經網路從簡單的前饋層序列快速演化為非常多樣的數值程式,通常由許多迴圈和遞迴函式組成。為了適應這一日益增長的複雜性,PyTorch 放棄了基於圖-超程式設計方法的潛在優勢,以保持 Python 的指令式程式設計模型。

PyTorch 將這一特點擴充套件到了深度學習工作流的所有方面。定義層、構建模型、載入資料、執行最佳化器、並行訓練過程都利用為通用程式設計開發的熟悉概念來表示。這一解決方案確保任何潛在的新神經網路架構都可以簡單地用 PyTorch 實現。

NeurIPS頂會接收,PyTorch官方論文首次曝光完整設計思路
一個簡單但完整的神經網路中用作構建塊的自定義層。

NeurIPS頂會接收,PyTorch官方論文首次曝光完整設計思路一個生成對抗網路的簡化訓練。

互操作性和可擴充套件性

PyTorch 允許與外部庫進行雙向交換。例如,PyTorch 提供了一種使用 torch.from_numpy() 函式和 .numpy() 張量方法的機制來實現NumPy 陣列和 PyTorch 張量使用之間的轉換。類似的功能也可用於交換使用 DLPack 格式儲存的資料。

此外,許多關鍵系統都是專門為可擴充套件性設計的。例如,自動微分系統允許使用者為自定義可微分函式新增支援。為此,使用者可以定義一個新的 torch.autograd.Function 子類來實現 forward() 和 backward() 方法,這些方法會指定函式及其導數。

自動微分

PyTorch 使用運算元超載(operator overloading)方法,在每次執行計算函式時構建一個該函式的表徵。在其最近的實現中,PyTorch 執行反向模式自動微分,計算有關多元輸入的標量輸出的梯度。

PyTorch 另一個有趣且不尋常的特性在於,它可以透過在張量上使用突變的程式碼進行微分,這是命令式程式的基本構建塊之一。

以效能為中心的實現

一個高效的 C++ 核

為了提高效能,PyTorch 的多數程式碼都是用 C++ 寫的。這一核心 libtorch 庫用來實現張量資料結構、GPU 和CPU 運算元以及基本的並行基元。它還提供了一個自動微分系統,包括用於多數內建函式的梯度公式。

Python 的 bingding 是使用 YAML 後設資料檔案生成的。這種方法的一個有趣副作用在於,它允許社群快速建立到多個其他語言的 binding ,產生了 NimTorch、hasktorch 等新專案。

分離控制和資料流

控制流的解由 Python 和最佳化的、在主機 CPU 上執行的 C++ 程式碼來處理,在裝置上產生一個運算元呼叫的線性序列。運算元可以在 CPU 或 GPU 上執行。

PyTorch 透過利用 CUDA 流機制將 CUDA 核心呼叫安排到 GPU 硬體 FIFO 來非同步執行運算元。

自定義快取張量分配器

PyTorch實現了一個自定義的分配器,它遞增地構建CUDA記憶體的快取並將其重新分配到之後的配額中,而無需進一步使用CUDA API。這種遞增的分配對於實現更好的互操作性也非常關鍵,因為提前佔用所有GPU記憶體會妨礙使用者利用其他GPU支援的Python包。為了進一步提高其效率,這一分配器針對深度學習的特定記憶體使用模式進行了調優。

這種「一流一池( one-pool-per-stream )」的設計假設簡化了實現,提高了分配器的效能。由於流序列化執行,如果空閒優先於 CPU 上的重新分配,同樣的順序也會發生在 GPU上。因此,只要在與釋放的區域相同的流上使用新的分配,分配器就可以立即重新分配在 CPU 上釋放的記憶體。

但這種設計似乎也是有限制的,因為每個流的分配結果是碎片化的,但在實際操作中,PyTorch 幾乎從不使用多個流。眾所周知,以一種讓CUDA 核心協同共享 GPU 的方式來編寫 CUDA 核心是非常困難的,因為精確的排程是由硬體控制的。

多程式處理

由於全域性直譯器鎖(global interpreter lock,GIL)的 Python 預設實現不允許並行執行緒進行並行執行,所以為了解決該問題,Python 社群已經建立了一個標準的多程式處理模組,其中包含了大量的實用程式(utility),它們可以使得使用者輕易地生成子程式並能夠實現基礎的程式間通訊原語(communication primitive)。

然而,原語的實現使用了與磁碟上永續性(on-disk persistence)相同格式的序列化,這在處理大規模陣列時效率不高。所以,PyTorch 將Python 的 multiprocessing 模組擴充套件為 torch.multiprocessing,這就替代了內建包,並且自動將傳送至其他程式的張量資料移動至共享記憶體中,而不用再透過通訊渠道傳送。

PyTorch 的這一設計極大地提升了效能,並且弱化了程式隔離(process isolation),從而產生了更類似於普通執行緒程式的程式設計模型。

引用計數

使用者常常設計模型來在訓練期間利用所有可用的記憶體,並且增加批次大小是加速程式的常見方法。所以,為了發揮出色的效能,PyTorch必須將記憶體視作稀有資源,並小心管理。

在引用計數方面,PyTorch 採取了一種不同的方法:它依賴於一個引用計數方案來追蹤每個張量的使用次數,並在該計數為零時立即釋放底層記憶體。需要注意的是,PyTorch 透過整合 Python 自身的引用機制,追蹤 libtorch 庫內部的引用以及使用者在其 Python 程式碼中所做的外部引用。

需要特別警醒的一點是,我們在已經利用引用計數的語言(CPython、Swift,而非 PyPy 或 Lua 等眾多指令碼語言)實現,或者在那些允許使用者自定義指定、複製和移動行為的語言(如 C++ 和 Rust )實現中只能保證預期的效能特性。


評估

研究者對 PyTorch 和其他幾個常用深度學習庫的效能進行了比較,發現 PyTorch 在一系列任務上都能實現較突出的效能。所有實驗都在一個使用兩個英特爾 Xeon E5-2698 v4 CPU 和一個英偉達 Quadro GP100 GPU 的工作站上執行。

非同步資料流

研究者首先量化了 PyTorch 在 GPU 上非同步執行資料流的能力。他們使用內建分析器來度量各種基準,並記錄下了單訓練步驟上的執行時間線。

下圖1展示了 ResNet-50 模型前幾步操作執行的典型時間線。在該例中,GPU 執行花費的時間約是 CPU 排程的3倍。精確的比例則取決於主 CPU 和 GPU 的相對效能、每個張量中的組成部件數量以及在 GPU 上實現的浮點運算的平均演算法複雜性。

NeurIPS頂會接收,PyTorch官方論文首次曝光完整設計思路

圖1: Resnet-50模型的前幾步操作的軌跡。

記憶體管理

研究者使用英偉達分析器來追蹤 CUDA 的執行時間以及 ResNet-50 模型訓練迭代期間啟動的 CUDA 核心的執行。如下圖2所示,首次迭代的行為表現與接下來的迭代截然不同。

NeurIPS頂會接收,PyTorch官方論文首次曝光完整設計思路圖2:GPU 上 ResNet-50 模型執行的自動追蹤。

基準測試

最後,透過與三個流行的圖深度學習框架(CNTK、MXNet 和 TensorFlow )、define-by-run 框架(Chainer)和生產導向型平臺(PaddlePaddle)的比較,研究者從整體上得出了 PyTorch 的單機 eager 模式效能。具體結果如下表1所示,PyTorch的效能在最快框架效能的17%以內。

NeurIPS頂會接收,PyTorch官方論文首次曝光完整設計思路表1:AlexNet、VGG-19、ResNet-50、MobileNet、GNMTv2 和 NCF 6 種模型在使用32位浮點運算時的訓練速度。

PyTorch 的應用

透過計算自 2017年1月PyTorch 首次釋出以來各種機器學習工具(包括 Caffe、Chainer、CNTK、Keras、MXNet、Pytorch、TensorFlow 和 Theano)在 arXiv 論文中提及的頻率,研究者試圖量化機器學習社群對 PyTorch 的接受程度。

如下圖3所示,研究者繪製出了所有這些深度學習框架中「PyTorch」每月出現的比例。

NeurIPS頂會接收,PyTorch官方論文首次曝光完整設計思路

圖3:自 2017 年1 月以來,在所有常見的深度學習框架中,PyTorch 在 arXiv 論文中每月被提及的比例。

未來展望

除了繼續支援深度學習領域最新的趨勢和進展之外,研究者計劃進一步提升 PyTorch 的速度和可擴充套件性。最主要的一點是,他們正開發 PyTorch JIT 系列工具,它們可以使得 PyTorch 程式脫離 Python 直譯器執行,從而可以得到進一步最佳化。研究者還打算透過為資料並行化提供高效的原語以及為基於遠端過程呼叫的模型並行化提供 Pythonic 庫,進而提升對分散式計算的支援。

相關文章