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

• 2020 年 3 月 27 日
• 筆記

機器之心編輯部

「深度學習，簡單開發」，剛剛開源的天元（MegEngine）將為開發者們帶來全新的深度學習開發體驗。

又一個國產深度學習框架開源了。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 已從實驗室中快速走向實用化。曠視的研發人員們認為，「天元」的開源宣告着人工智能技術再次回到了算法探索的層面。

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

本文為機器之心報道，轉載請聯繫本公眾號獲得授權。