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

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

1. 用戶在 python 層編寫網(wǎng)絡(luò)結(jié)構(gòu)代碼，執(zhí)行時向 C++ 層發(fā)射算子執(zhí)行指令
2. Imperative 的 dispatcher 對部分算子做計算前的預(yù)處理（transformation）
3. Imperative 的 interpreter 執(zhí)行計算操作（復(fù)用 MegBrain 的相關(guān)組件）

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

MegEngine 的 Python 層

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

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

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

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

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

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

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

Python 層計算接口 —— functional 模塊

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

functional 中實現(xiàn)了各類計算函數(shù)，包含對很多 op 的封裝，供實現(xiàn)模型時調(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)或者層，用戶實現(xiàn)算法時往往使用組合 Module 的方式搭建模型計算的 pipeline。定義神經(jīng)網(wǎng)絡(luò)時有些結(jié)構(gòu)經(jīng)常在模型中反復(fù)使用，將這樣的結(jié)構(gòu)封裝為一個 Module，既可以減少重復(fù)代碼也降低了復(fù)雜模型編碼的難度。

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

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

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

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

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

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

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

Dispatcher 會對 op 做哪些處理？

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

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

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

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

• DimExpansionTransformation：某些 op 計算時對輸入 tensor 的 shape 有要求，在這里做處理。

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

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

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

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

• TracingTransformation：

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

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

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

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

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

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

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

因為不同的 Transformation 有邏輯上的先后關(guān)系，所以開發(fā)者往往需要手動規(guī)劃它們之間的順序。

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

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

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

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

這里我們需要先了解一個 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 的接口需要實現(xiàn)對應(yīng)的方法，在 MegEngine/imperative/src/include/megbrain/imperative 目錄下可以看到所有需要實現(xiàn)的橋接接口，其中和 proxy_graph 相關(guān)的接口聲明在 proxy_graph_detail.h 中，通常需要實現(xiàn)這幾個接口：

• infer_output_attrs_fallible

  復(fù)用 MegBrain 的 StaticInferManager 進(jìn)行 shape 推導(dǎo)，在執(zhí)行計算操作前對輸入和輸出 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í)行計算操作。

• make_backward_graph

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

• get_input_layout_constraint

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

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

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

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