在之前的文章中我們介紹過?MegEngine?的?Imperative Runtime?以及它與?MegBrain、MegDNN?的關(guān)系，這篇文章中我們將介紹?Imperative?中包含的常用組件。

在?MegEngine?中，從用戶在?python?層編寫代碼到在?interpreter?層發(fā)生計(jì)算經(jīng)過了下面的流程：

1. 用戶在?python?層編寫網(wǎng)絡(luò)結(jié)構(gòu)代碼，執(zhí)行時(shí)向?C++?層發(fā)射算子執(zhí)行指令

2. Imperative?的?dispatcher?對部分算子做計(jì)算前的預(yù)處理（transformation）

3. Imperative?的?interpreter?執(zhí)行計(jì)算操作（復(fù)用?MegBrain?的相關(guān)組件）

我們將分別介紹這幾個(gè)階段系統(tǒng)所做的工作。

MegEngine 的 Python 層

在主流的深度學(xué)習(xí)框架中，用戶往往不需要自己手寫算子的具體實(shí)現(xiàn)、處理計(jì)算圖的執(zhí)行邏輯、或者與復(fù)雜的體系結(jié)構(gòu)打交道。一切都被封裝為?Python?層的接口。

在?MegEngine?的?Python?層中用戶接觸較多的模塊主要有：data、functional、module、optimizer、quantization、tools，下面簡單介紹一下各個(gè)模塊的功能。

構(gòu)建數(shù)據(jù)處理 Pipeline —— data 模塊

Data?模塊，顧名思義就是對數(shù)據(jù)進(jìn)行處理的模塊。

沒有數(shù)據(jù)就沒法訓(xùn)練，在?MegEngine?中，通常會借助一個(gè)?Dataset?結(jié)構(gòu)來定義數(shù)據(jù)集。數(shù)據(jù)集一般分為?Map-stype?和?Iterable-style?兩種。前者叫作?ArrayDataset，這種數(shù)據(jù)集支持隨機(jī)訪問；后者叫作?StreamDataset，因?yàn)槭橇魇降臄?shù)據(jù)集，只支持順序訪問。

有了數(shù)據(jù)集，我們還需要一個(gè)結(jié)構(gòu)來把數(shù)據(jù)“喂”給模型訓(xùn)練，這樣的一個(gè)結(jié)構(gòu)叫作?dataloader。

實(shí)際上，只給?dataloader?一個(gè)?dataset?有時(shí)無法準(zhǔn)確地描述加載數(shù)據(jù)的整個(gè)過程，我們可能還需要定義加載數(shù)據(jù)過程的抽樣規(guī)則（Sampler），或者定義一些數(shù)據(jù)變換的規(guī)則（Transform），或者是定義抽樣后的數(shù)據(jù)的合并策略（Collator）。

Python 層計(jì)算接口 —— functional 模塊

深度學(xué)習(xí)模型通常包含一些基礎(chǔ)的計(jì)算操作，比如?convolution、pooling?等，在 python 層，這些基本計(jì)算操作都定義在?functional?模塊中。

functional?中實(shí)現(xiàn)了各類計(jì)算函數(shù)，包含對很多?op?的封裝，供實(shí)現(xiàn)模型時(shí)調(diào)用。

模型結(jié)構(gòu)的小型封裝版本 —— module 模塊

使用?functional?提供的接口已經(jīng)足夠編寫神經(jīng)網(wǎng)絡(luò)模型的代碼，但隨著模型結(jié)構(gòu)的復(fù)雜程度加深，多次反復(fù)編寫相似的結(jié)構(gòu)會使開發(fā)和維護(hù)成本迅速提高。

考慮到神經(jīng)網(wǎng)絡(luò)模型通常是由各種層（layer）組成，我們通常使用??Module? 來封裝模型的部分結(jié)構(gòu)或者層，用戶實(shí)現(xiàn)算法時(shí)往往使用組合??Module? 的方式搭建模型計(jì)算的?pipeline。定義神經(jīng)網(wǎng)絡(luò)時(shí)有些結(jié)構(gòu)經(jīng)常在模型中反復(fù)使用，將這樣的結(jié)構(gòu)封裝為一個(gè)?Module，既可以減少重復(fù)代碼也降低了復(fù)雜模型編碼的難度。

使用 optimizer 模塊優(yōu)化參數(shù)

MegEngine?的?optimizer?模塊中實(shí)現(xiàn)了大量的優(yōu)化算法，同時(shí)為用戶提供了包括?SGD、?Adam?在內(nèi)的常見優(yōu)化器實(shí)現(xiàn)。 這些優(yōu)化器能夠基于參數(shù)的梯度信息，按照算法所定義的策略對參數(shù)執(zhí)行更新。

降低模型內(nèi)存占用利器 —— quantization 模塊

量化是一種對深度學(xué)習(xí)模型參數(shù)進(jìn)行壓縮以降低計(jì)算量的技術(shù)。它基于這樣一種思想：神經(jīng)網(wǎng)絡(luò)是一個(gè)近似計(jì)算模型，不需要其中每個(gè)計(jì)算過程的絕對的精確。因此在某些情況下可以把需要較多比特存儲的模型參數(shù)轉(zhuǎn)為使用較少比特存儲，而不影響模型的精度。

MegEngine 相關(guān)工具匯總 —— tools 模塊

用戶進(jìn)行開發(fā)時(shí)有時(shí)需要一些工具進(jìn)行錯(cuò)誤調(diào)試或者性能調(diào)優(yōu)，tools?下就提供了一些這樣的工具。比如對訓(xùn)練程序進(jìn)行記錄并在瀏覽器上可視化的?profiler、方便用戶查看?MegEngine?顯存占用的?svg_viewer?等。

一般來說，用戶會基于上面的模塊搭建算法模型，其中定義了非常多的 op 的計(jì)算過程，下面我們看一下 c++ 是怎么進(jìn)行這些 op 的真正的計(jì)算的。

Dispatcher 會對 op 做哪些處理？

從?Python?層往下的部分用戶往往是感知不到的，脫離了“前端”，我們抽絲剝繭，進(jìn)入到了框架“后端”對?tensor?和?op?處理的細(xì)節(jié)。

前面我們提到在?functional?模塊中封裝了很多算子，并以?python?接口的形式提供。實(shí)際上這些算子需要向下發(fā)射指令對?tensor?進(jìn)行操作并返回操作完成后的?tensor，這些發(fā)射的?op?指令就會到?dispatch?層，在進(jìn)行實(shí)際計(jì)算之前，dispatcher?會對?tensor?做一些處理，我們把這些處理叫作?Transformation。

在?imperative?中真正執(zhí)行算子進(jìn)行計(jì)算是在?interpreter?層做的，與?tensor?處理相關(guān)的操作被解耦出來放在?dispatch?層，這樣更便于維護(hù)。

在?MegEngine?中，一些重要的?transformation?有：

• DimExpansionTransformation：某些?op?計(jì)算時(shí)對輸入?tensor?的?shape?有要求，在這里做處理。

• DtypePromoteTransformation：某些?op?要求計(jì)算的?tensor?擁有相同的類型，會將所有的輸入的類型提升為同一類型之后再進(jìn)行計(jì)算。比如?int?類型?tensor?和?float?類型?tensor?進(jìn)行計(jì)算，需要把?int?類型的?tensor?轉(zhuǎn)換為?float?類型?tensor。

• InterpreterTransformation：顧名思義，這類?Transformation?將指令轉(zhuǎn)發(fā)到?Interpreter?層（Interpreter?可以認(rèn)為是?Imperative?中所有計(jì)算操作的入口）進(jìn)行計(jì)算，并獲取指令的計(jì)算結(jié)果。Transformation?通常是疊加的，InterpreterTransformation?是最后一層，其后不再跟其他的?Transformation?處理。

• FormatTransformation：由于在不同情況下對不同?format?的?Tensor?的計(jì)算速度不同，因此需要對?NHWC?和?NCHW?的?Tensor?進(jìn)行轉(zhuǎn)換，為了不讓用戶感知到這樣的轉(zhuǎn)換，這部分的工作由?FormatTransformation?完成。

• GradTransformation：訓(xùn)練模型時(shí)需要通過反向傳播更新模型參數(shù)，反向傳播需要支持?op?的自動(dòng)微分。要實(shí)現(xiàn)求導(dǎo)，就需要在前向執(zhí)行?op?的時(shí)候記錄某些信息，以便之后進(jìn)行反向求導(dǎo)。Autodiff?算法會根據(jù)輸入的前向圖生成一個(gè)完整的前向反向圖，所謂的前傳反傳訓(xùn)練過程對?Autodiff?來說實(shí)際上都是一個(gè)計(jì)算圖的前向過程，grad?的數(shù)值是在“前向”的過程中就已經(jīng)拿到的。GradTransformation?處理的就是與反向求導(dǎo)相關(guān)的操作。

• TracingTransformation：

  在介紹?Trace?之前，我們需要先明確一下計(jì)算圖的概念。計(jì)算圖可以認(rèn)為是對輸入的數(shù)據(jù)（tensor）、op?以及?op?執(zhí)行的順序的表示。計(jì)算圖分為動(dòng)態(tài)圖和靜態(tài)圖。動(dòng)態(tài)圖是在前向過程中創(chuàng)建、反向過程銷毀的。前向邏輯本身是可變的，所以執(zhí)行流程也是可變的（因此叫動(dòng)態(tài)圖），而靜態(tài)圖的執(zhí)行流程是固定的。也就是說，動(dòng)態(tài)圖在底層是沒有嚴(yán)格的圖的概念的（或者說這個(gè)圖本身一直隨執(zhí)行流程變化）。對于動(dòng)態(tài)圖來說，graph?的?node?對應(yīng)的概念是?function?/ 算子，而?edge?對應(yīng)的概念是?tensor，所以在圖中需要記錄的是?graph?中?node?和?edge?之間的連接關(guān)系，以及?tensor?是?function?的第幾個(gè)輸入?yún)?shù)。

  Trace?的作用就是將動(dòng)態(tài)圖執(zhí)行轉(zhuǎn)換為靜態(tài)圖執(zhí)行，這樣做的好處就是執(zhí)行速度更快了，并且占用的顯存更少了。因?yàn)殪o態(tài)圖需要先構(gòu)建再運(yùn)行，可以在運(yùn)行前對圖結(jié)構(gòu)進(jìn)行優(yōu)化（融合算子、常數(shù)折疊等），而且只需要構(gòu)建一次（除非圖結(jié)構(gòu)發(fā)生變化）。而動(dòng)態(tài)圖是在運(yùn)行時(shí)構(gòu)建的，既不好優(yōu)化還會占用較多顯存。

  Trace?中所有的東西都會進(jìn)行靜態(tài)優(yōu)化（加速）。

  加了?Trace?之后，模型在訓(xùn)練時(shí)第一個(gè)?iter?是動(dòng)態(tài)圖執(zhí)行，Trace?會記錄下?tensor、op?以及?op?的執(zhí)行順序這些信息（構(gòu)建靜態(tài)圖）并進(jìn)行計(jì)算，在第二個(gè)?iter?就跑的是構(gòu)建好的靜態(tài)圖。

• LazyEvalTransformation：類似?TracingTransformation，也會記錄?tensor、op?等信息構(gòu)建靜態(tài)圖，不同的是?LazyEvalTransformation?在第一個(gè)?iter?不會跑動(dòng)態(tài)圖，但會在第二個(gè)?iter?開始跑靜態(tài)圖。

• ScalarTransformation：用于判斷指令的輸出是否為?scalar。因?yàn)?dispatch?的?Tensor?要發(fā)到?Interpreter?層，而?Interpreter?層不接受?ndim == 0?的?Tensor（在?Interpreter?中?ndim?為?0?表示?Tensor?的?shape?未知），也就是一個(gè)?scalar，因此?ScalarTransformation?會將?ndim?為?0?的?Tensor?表示為?ndim?不為?0?的?Tensor?（具體是多少與具體?op?有關(guān)）發(fā)往?Interpreter。

不同的?Transformation?之間擁有固定的執(zhí)行順序：比如?InterpreterTransformation?是執(zhí)行實(shí)際計(jì)算并獲取計(jì)算結(jié)果的（需要進(jìn)入?Interpreter），所以它是在最后一個(gè)執(zhí)行的。TracingTransformation?/?LazyEvalTransformation?/?CompiledTransformation?等屬于?Trace?相關(guān)的操作，因?yàn)?Trace?需要記錄所有指令，所以這些?Transformation?是在倒數(shù)第二層執(zhí)行的。如?ScalarTransformation?這樣只對?Scalar?做處理的?Transformation?往往在較上層。

因?yàn)椴煌?Transformation?有邏輯上的先后關(guān)系，所以開發(fā)者往往需要手動(dòng)規(guī)劃它們之間的順序。

不同類型的?Transformation?之間是解耦的，這樣便于開發(fā)與維護(hù)。

Interpreter 是如何“解釋”算子的？

由于?MegBrain?已經(jīng)是一個(gè)非常成熟的靜態(tài)圖框架，因此在開發(fā)動(dòng)態(tài)圖（Imperative Runtime）深度學(xué)習(xí)框架?MegEngine?的過程中，復(fù)用許多靜態(tài)圖中的組件可以大大降低開發(fā)成本。

實(shí)際上，張量解釋器?Tensor Interpreter?就是將動(dòng)態(tài)圖中的操作——如執(zhí)行?op、shape?推導(dǎo)等操作“解釋”為靜態(tài)圖的對應(yīng)操作，并復(fù)用?MegBrain?的組件來運(yùn)行。

這里我們需要先了解一個(gè)?MegBrain?的靜態(tài)圖“長什么樣”。

復(fù)用靜態(tài)圖接口的機(jī)制 —— proxy_graph

為了復(fù)用?MegBrain?的靜態(tài)求導(dǎo)器、靜態(tài)內(nèi)存分配器、靜態(tài)?shape?推導(dǎo)器等組件，imperative?引入了?proxy_graph。

復(fù)用?MegBrain?的接口需要實(shí)現(xiàn)對應(yīng)的方法，在?MegEngine/imperative/src/include/megbrain/imperative?目錄下可以看到所有需要實(shí)現(xiàn)的橋接接口，其中和?proxy_graph?相關(guān)的接口聲明在?proxy_graph_detail.h?中，通常需要實(shí)現(xiàn)這幾個(gè)接口：

• infer_output_attrs_fallible

  復(fù)用?MegBrain?的?StaticInferManager?進(jìn)行?shape?推導(dǎo)，在執(zhí)行計(jì)算操作前對輸入和輸出?tensor?的?shape?進(jìn)行檢查。

• apply_on_physical_tensor

  根據(jù)?infer_output_attrs_fallible?推導(dǎo)的?shape?結(jié)果去分配?op?輸出的顯存，并調(diào)用?proxy opr?的?execute?函數(shù)（會轉(zhuǎn)發(fā)到?MegDNN?的?exec?函數(shù)）執(zhí)行計(jì)算操作。

• make_backward_graph

  在求導(dǎo)時(shí)，Grad Manager?會記錄下來一些求導(dǎo)需要的信息（輸入?tensor、op?以及它們執(zhí)行的順序、輸出?tensor），make_backward_graph?會根據(jù)這些信息造一個(gè)反向的計(jì)算圖，供求導(dǎo)使用。

• get_input_layout_constraint

  一般用來判斷一個(gè)輸入?tensor?的?layout?是否滿足一些限制：比如判斷?tensor?是否是連續(xù)的。

  如果不滿足限制，則會造一個(gè)滿足限制的?tensor，供?apply_on_physical_tensor?使用。

在實(shí)現(xiàn)一個(gè)?imperative?算子時(shí)通常也只需要實(shí)現(xiàn)這幾個(gè)接口，剩下的工作由?MegBrain?和?MegDNN?完成。

主流框架在?python?層的模塊封裝結(jié)構(gòu)大同小異，關(guān)于?MegEngine?的?Python?層各模塊的使用與實(shí)現(xiàn)細(xì)節(jié)以及?transformation?和?interpreter?實(shí)現(xiàn)細(xì)節(jié)我們會在之后的文章中逐一解析。

附

更多 MegEngine 信息獲取，您可以：查看文檔：https://www.megengine.org.cn/doc/stable/zh/

GitHub 項(xiàng)目： https://github.com/MegEngine

加入 MegEngine 用戶交流 QQ 群：1029741705

歡迎參與 MegEngine 社區(qū)貢獻(xiàn)，成為?Awesome MegEngineer：https://www.megengine.org.cn/community-AMGE，榮譽(yù)證書、定制禮品享不停。