訓推一體，動靜合一：深度學習框架「天元」出世

機器之心發表於2020-03-25
深度學習框架

又一個國產深度學習框架開源了。3 月 25 日，曠視科技在北京釋出了新一代人工智慧生產平臺天元（MegEngine），其「訓練推理一體」、「動靜合一」、「相容幷包」、「靈活高效」的強大效能或許將為 AI 領域帶來新的風向。

「天元是曠視 1400 名研發人員一直以來使用的深度學習引擎，已部署在上百個產品、幾十種計算平臺上，」曠視首席科學家、研究院長孫劍在釋出會上表示。「他可以幫助開發者做到協同設計訓得好，高效系統訓得快，大規模資料或模型訓得動。」

訓推一體，動靜合一：深度學習框架「天元」出世

曠視首席科學家、研究院長孫劍釋出會現場分享。

曠視提出的 AI 生產力平臺 Brain++包括深度學習框架（天元 MegEngine）、深度學習雲端計算平臺（MegCompute）、以及資料管理平臺（MegData），將演算法、算力和資料能力集為一體。本次開源的天元是 Brain++平臺的核心。

  • 天元 MegEngine 深度學習框架官方網站：https://megengine.org.cn/

  • GitHub 開源地址：https://github.com/MegEngine/MegEngine


時至 2020，自賈揚清等人開放 Caffe 已過去六年，曠視為什麼選擇在這個時間點開源自身的核心技術？曠視研究院高階技術總監、天元專案負責人田忠博告訴我們：「這一程式其實去年已經開始推動，我們認為目前國內市場上缺少天元這樣的專案。」

這是一個需要勇氣的行動。據瞭解，在做出開源的決定以前，曠視技術團隊內部經過了多次討論。這關係到框架技術是否成熟，生態環境是否發展完善，以及未來發展的方向。

田忠博指出，天元是一套訓練推理一體化、動靜態合一的工業級深度學習框架，具體架構包括五層：

訓推一體，動靜合一：深度學習框架「天元」出世


從最上層的介面層到最底層的計算核心，天元做了很多優化，例如記憶體管理中的亞線性優化和高效能運算元庫等。

天元架構是曠視在研發過程中遇到痛點，並逐步解決的成果。曠視相信這些痛點也是行業共通的難題，通過這一開源架構，我們可以一站式地解決這些問題。

早在 2014 年 8 月，深度學習框架天元就在曠視內部誕生了，它隨後成為了 Brain++的核心組成部分。「自 2012 年起，陸續有研究者開發出了第一代深度學習框架。在曠視，我們在實踐過程中發現手工編寫每層配置檔案的方式非常不靈活，」孫劍介紹道。「因此我們發展了第二代深度學習框架，通過計算圖的方式設計模型並求導，大大提高了效率。2015 年穀歌正式釋出了 TensorFlow，採用了和曠視殊途同歸的思路。當時我們進行了比較，發現自己的方法更具優勢，所以天元就繼續發展下來了。」

天元與目前流行的 TensorFlow、PyTorch 相似，旨在讓更多的人可以在其基礎上構建自己的應用，創造屬於自己的價值。

此前，百度等公司也推出了自己的深度學習框架。天元會有哪些不同？「在這其中，我們是唯一一個以 AI 為本業的公司，天然需要使用深度學習演算法解決問題，」曠視副總裁謝憶楠表示，「我們把自己在 AI 解決各領域問題過程中的工具開源出來，從視角上看或許會獲得更好的效果。」

天元躋身全球主流框架

機器之心很多讀者都用過深度學習框架，很多讀者也都想入門深度學習框架。那麼假設我們有 Python 和 NumPy 數值計算方面的基礎，曠視新開源的天元又會有哪些方面吸引我們？

孫劍表示，天元有三大特別吸引人的優點，「第一是框架與演算法的協同，框架需要為硬體開發優化運算元，這樣才能結合框架、裝置開發最好的網路，ShuffleNet 系列高效神經網路就是協同設計的結果之一。第二即高效與強大效能，一般訓練 COCO 資料集需要 30 多個小時，2017 年基於天元的 MegDet 並行物體檢測系統，2 小時就能完成，並且精度更高。」

「第三是很好地支援大資料、大規模訓練。借天元開源之際，我們同時也釋出最新版 Objects365 V2，它是世界上最大的物體檢測資料集，是目前學界流行的 COCO 資料集 的 16 倍大小。有這麼大規模的資料，模型也需要擴增。如果模型擴大 10 倍的話，總體訓練量就會增加 160 倍。而天元對於這種大模型、大資料也是沒問題的，」孫劍接著說。

訓推一體，動靜合一：深度學習框架「天元」出世

田忠博說：「天元是為工業級、研究院規模的研發機構設計的，它具備的一些核心特點都是用來解決現實研發過程中的痛點。」從具體體驗與特性上，我們確實能發現天元解決很多開發痛點，它們從程式碼風格、程式設計正規化到底層優化展現了不一樣的思路。

一鍵切換訓練推理

在機器學習中，訓練與推理一直不是一體的，當我們設計演算法時，首先需要訓練框架的支援，然後再把訓練好的模型轉移到推理框架上，從而將後續的推理計算放到各種硬體上。這中間會存在訓練與推理的轉換，這個過程需要手動完成，甚至轉換後也不能保證模型效能與進度。

「訓練推理一體化能把繁雜的模型生產流程簡化到一步、兩步，我們也希望它能夠極大縮短演算法研發和部署的時間，真正支援大規模的生產級訴求」，田忠博說，「為此，天元框架需要做到三點：無需模型轉換就能做推理；訓練的速度、精度與推理保持一致；在不同裝置上做推理需要保證模型精度可以對齊。」

訓推一體，動靜合一：深度學習框架「天元」出世

從演算法直接到部署硬體，中間不再需要轉換模型，只需要幾行程式碼與 API，模型的實踐效果就能看得到。對於大多數演算法工程師，可能比較精通深度學習模型，對推理與模型部署相對而言不是那麼熟，天元無疑能降低大量工程實踐成本。

更重要的是，在部署過程中，天元通過自動演算法選擇機制和一系列針對推理的自動優化，能幫我們刪除冗餘程式碼，找到最適宜該裝置的演算法實現，達到靈活高效的目標。

動靜合一，圖並存

從 Theano 到 TensorFlow 1.X，它們都遵循著「靜態計算圖」這條路，我們會發現在用這類方法寫模型時，需要很強的「向量化思維」，For 迴圈之類的語句就不要多想了。此外，因為寫程式碼只是建計算圖，除錯的時候異常困難，我們不知道圖中的資料流到底是什麼樣的。

很快，在 PyTorch 等框架中出現了動態計算圖，用它寫模型就像寫個普通 Python 指令碼，但這樣寫出來的模型，卻是難以高效部署的。天元同時支援動態圖與靜態圖，在科研訓練與除錯時採用動態圖，到部署的時候又切換到靜態圖，這樣就能充分利用兩者優勢。

訓推一體，動靜合一：深度學習框架「天元」出世

如上圖所示，天元可以直接通過 @trace 裝飾器確定某個函式到底是動態圖還是靜態圖，而且狀態還可以手動調整。這不就意味著在除錯時我們可以將函式設定為動態圖，直接檢視每個變數的資訊，而要到了部署等實際應用時，再把靜態圖開啟，提升訓練與推斷速度？

如果以這樣的方式呼叫動態與靜態計算圖，那真的稱得上整合了兩種方式的優勢。

此外，我們還發現有一種動、靜態的「混合程式設計」，即在動態圖中呼叫靜態圖。可以想象，定義網路結構等過程使用動態圖比較直觀，而後採用靜態圖的方式執行實際運算，這也是非常吸引人的方式。後文在程式碼示例中，也會展示這種新穎的「混合程式設計」。

人生苦短，Pythonic 好入門

如果深度學習建模一直採用靜態計算圖，估計框架就會勸退很多開發者。但是天元這種主要採用動態圖的框架，構建模型就像搭積木，而且還能以一般的 Python 程式設計方式去搭。相信很多接觸機器學習的開發者都會選擇 Python，它加上 NumPy 等數值計算庫，差不多就能完成大多數任務。

如果有了這些基礎，再搭建深度模型就沒問題了，天元一樣承接 Pythonic API，學習成本非常低。此外，因為曠視本身在計算機視覺方面有非常雄厚的累積，田忠博表示：「天元內建了一些獨特的深度學習運算元，它們和曠視在 CV 上的獨到能力息息相關，能針對 CV 模型的訓練與應用實現深度優化。」

訓推一體，動靜合一：深度學習框架「天元」出世

對於 API 與程式碼風格，我們還是比較推薦直接檢視本文後面的示例，可以直觀感受到，不論把天元當做深度學習的入門框架，還是與 PyTorch、TensorFlow 等一起作為備選框架都沒問題。天元的程式設計風格與 API 配置，和已有的程式設計習慣、現有框架的風格都是一致的。

當然，這些並不是天元所有特性，還有很多需要進一步探索。例如它還有很強的多平臺、多裝置適應能力，既能靈活呼叫硬體能力，又可以高效訓練大模型。

天元框架實際體驗怎麼樣？

那麼實際用天元寫模型是什麼樣的體驗？機器之心體驗了曠視釋出的 MegStudio，通過線上開發的方式，我們可以快體驗到天元的特性與程式設計方式。

MegStudio 上準備好了一些入門專案，它展示了新框架的使用方式，例如「基本概念」主要介紹什麼是計算圖、張量和運算元等基礎知識，每一個專案都是通過中文文字與可執行的程式碼組成的，就像常用的 JupyterLab 那樣，通過執行這些基本專案可以讓開發者快速掌握天元的基本使用。

訓推一體，動靜合一：深度學習框架「天元」出世

為了直接展示天元的風格，我們可以快速試試 DL 入門的「第一個模型」MNIST 影像識別。準備資料這裡就不展示了，天元可以直接呼叫 API 來處理常見資料，準備 MNIST 資料兩行程式碼就夠了。如果要搭建一個簡單的卷積神經網路，天元允許使用 PyTorch Module 的方式定義，搭建模型就是搭積木。

訓推一體，動靜合一：深度學習框架「天元」出世

如上所示，繼承 M.Module 後，Net 類只需要定義初始化引數與重寫前饋 forward 過程就可以了。卷積層、批歸一化層、啟用函式等模組的使用非常符合直接，只需要給定必要引數就沒問題了。

這種繼承 Module 的方式非常直觀，是一種動態計算圖編寫方式。如果要實現前向傳播，傳給它輸入資料就行了，要是需要除錯，在 forward 函式下設定斷點，或者列印變數的資訊都沒問題。

然而，一般寫模型真正複雜的地方在於定義模型架構之外的過程，例如後面需要完成的訓練與測試過程。天元的優勢在於，它允許動態圖結合靜態圖，即使定義模型使用的是動態圖，但在定義訓練與測試函式仍然可以用靜態計算圖。

訓推一體，動靜合一：深度學習框架「天元」出世


如上在最為核心的前向、反向傳播與模型損失計算過程中，我們可以採用天元中的 @trace 裝飾器，直接呼叫靜態計算圖模式。因為絕大多數計算都發生在這個過程，因此，由動態計算圖轉為靜態計算圖將更有優勢，天元能採用更高效的方式訓練模型。

最後就是迭代訓練過程了，這裡天元也採用了 DataLoader，這種方式非常便捷。後續的損失函式計算與最優化器更新，與現有的 DL 框架比較類似，我們差不多不需要學習額外的知識，瞭解一下 API 就能用。

訓推一體，動靜合一：深度學習框架「天元」出世


當然，後續還有儲存、載入模型，也是兩條語句的功夫。總體而言，很明顯天元類似 PyTorch 保持了最好理解與使用的介面，並在關鍵的計算過程提供裝飾器以呼叫靜態計算圖，這可以稱得上是結合了動態、靜態圖各自的優勢。

因為整體上天元的 API 介面非常精簡，且保持著主流框架所養成的程式設計習慣，入門門檻還是非常低的，之前瞭解過深度學習框架的同學，不需要文件都能快速讀懂整個過程。此外天元作為新手入門框架也沒是沒問題的，學會了它，差不多等同於學會了深度學習常用框架。

此外，除了 MegStudio，曠視還發布了模型中心 ModelHub，它匯聚了全球頂尖演算法的一些預訓練的模型，和曠視研究院一些最新的技術、研發成果。

訓推一體，動靜合一：深度學習框架「天元」出世

曠視聯合創始人兼 CTO 唐文斌 釋出會現場分享

在曠視的計劃中，Brain++將作為「AI 基礎設施」逐步開源、開放，將生產力共享給廣大開發者，讓每個擁有小場景需求的人都可以做出能夠解決其關鍵問題的最優演算法。此外，曠視還希望在未來面向高校課程、初創企業提供更多支援，以方便小規模深度學習部署。

值得一提的是，曠視本次釋出的天元版本是 Alpha 版。曠視表示，更為「正式」的版本會在今年晚些時候上線，隨著大量元件的更新，曠視會逐步將計劃中開源的程式碼和能力進一步釋放出來。「3 月份開源的版本包含前期整理的程式碼和關鍵元件，今年 6 月我們還會發布對 Arm 架構、深度學習加速晶片、低位元和量化環境的支援。在今年 9 月份，天元還會有動態圖的核心升級，整個系統會走向第一個開源的穩定版。屆時體驗會達到更為理想的狀態。」田忠博表示。

自 2012 年起，曠視就推出了人臉識別開放平臺 Face++，近幾年來，AI 已從實驗室中快速走向實用化。曠視的研發人員們認為，「天元」的開源宣告著人工智慧技術再次回到了演算法探索的層面。

未來，曠視還會進一步增強這一框架的分散式機器學習能力，對底層編譯與程式碼生成能力進行優化，並歡迎機器學習社群力量的加入。

相關文章