TracedModule 介紹

TracedModule 是 MegEngine 中的一種模型格式，用于脫離模型源碼對(duì)模型進(jìn)行訓(xùn)練、量化、圖手術(shù)和模型轉(zhuǎn)換，它是模型訓(xùn)練到部署之間的橋梁。?

TracedModule 產(chǎn)生自普通的?Module，它通過以下兩步得到：

• 運(yùn)行一次 Module，記錄并捕獲模型運(yùn)行過程中對(duì)輸入 Tensor 的所有操作，對(duì)應(yīng)圖 1 中的?tm.trace_module

• 通過一個(gè)由 5 種指令所構(gòu)成的 “簡(jiǎn)單” 的 high-level IR(intermediate representation) 來描述捕獲到的程序（普通 Module 中的 forward 方法），對(duì)應(yīng)于圖 1 中的?SimpleModule.Graph

TracedModule 的本質(zhì)仍然是一個(gè) Module，它與普通 Module 的區(qū)別在于： 普通 Module 通過用戶實(shí)現(xiàn)的 forward 方法描述模型運(yùn)行過程，而 TracedModule 通過 TracedModule IR 來描述模型的運(yùn)行過程。TracedModule IR 由 python 的基本數(shù)據(jù)類型以及 Node 和 Expr 構(gòu)成，其中：Node 用來表示一個(gè) Tensor 或 Module，并記錄了 Tensor 或 Module 的一些信息；Expr 用來表示對(duì) Tensor 或 Module 的操作，它的輸入和輸出都是 Node。

TracedMdoule IR 中的 Expr 共有以下 5 種：

通過以上 5 種 Expr 即可表示絕大部分模型的運(yùn)行過程。

為什么要有 TracedModule？

如前文的介紹，TracedModule IR 是 TracedModule 中的核心數(shù)據(jù)結(jié)構(gòu)，它用來描述深度學(xué)習(xí)模型的計(jì)算過程，使模型能夠脫離源碼而存在。不同的深度學(xué)習(xí)訓(xùn)練框架都有各自的 IR 描述模型，例如：PyTorch 中的 TorchScript，MindSpore 中的 MindIR，以及 onnx 等。這些 IR 大都是一些相對(duì) low-level 的 IR，在模型源碼向這些 IR 轉(zhuǎn)換時(shí)常常會(huì)發(fā)生 python 的層 op 被轉(zhuǎn)換為多個(gè)框架底層 op 組合的現(xiàn)象，例如 pytorch 中的?F.Linear?算子在導(dǎo)出到 TorchScript 時(shí)可能被導(dǎo)出為?matmul?和?add?的組合。用戶在使用 low-levle 的 IR 表達(dá)的模型時(shí)會(huì)有很多的問題，例如：

• 不了解底層算子用戶可能會(huì)很難從模型的可視化結(jié)構(gòu)上與模型源碼對(duì)應(yīng)

• 普通用戶學(xué)習(xí) IR 結(jié)構(gòu)較為困難，很難對(duì)模型進(jìn)行圖手術(shù)（修改模型圖結(jié)構(gòu)）或優(yōu)化

這種 low-level 的 IR 表達(dá)能力往往更加完備，相應(yīng)的也導(dǎo)致 IR 結(jié)構(gòu)極其復(fù)雜，失去高層語(yǔ)義，使得用戶難以做變換和優(yōu)化，對(duì)模型設(shè)計(jì)者非常不友好。業(yè)界也提出了一些更加 high-level 的 IR 來解決這些問題，比如 torch.fx 和 pnnx 等，這些 IR 都對(duì)結(jié)構(gòu)進(jìn)行了簡(jiǎn)化，讓 IR 的描述模型中的 op 粒度更高，更貼近算法工程師的視角，使得用戶學(xué)習(xí)更簡(jiǎn)單，處理模型也更容易。

在 MegEngine 中，由多個(gè)底層 op 組合成的 python 層 op 更多，例如 "resize", "relu6", "softmax" 等，如果直接通過底層 op 表達(dá)模型，將會(huì)出現(xiàn)導(dǎo)出的模型結(jié)構(gòu)誰(shuí)也不認(rèn)識(shí)的窘?jīng)r。為了解決這些問題，MegEngine 參考 torch.fx 和 TorchScript 方案，改進(jìn)得到 TracedModule 方案。TracedModule 的 IR 是一個(gè) high-level 的 IR，它所描述的 op 粒度基本與 MegEngine 的 python 層的 op 一致，模型中的 op 粒度與用戶視角一致，用戶可以很容易地基于 TracedMdoule IR 對(duì)模型進(jìn)行分析，優(yōu)化和轉(zhuǎn)換。另外前文提到 TracedMdoule 的本質(zhì)是一個(gè) Module，用戶也可以方便地使用 MegEngine 的模型訓(xùn)練接口對(duì) TracedModule 模型進(jìn)行訓(xùn)練或參數(shù)微調(diào)。

TracedModule 好在哪？

TracedModule 全部由 python 層的數(shù)據(jù)結(jié)構(gòu)構(gòu)成，trace_module?函數(shù)在捕獲用戶代碼的運(yùn)行邏輯時(shí)僅記錄模型中使用的 MegEngine python 層的 function 或 Module，這使得 TracedModule IR 所描述的 op 粒度基本與 MegEngine 的 python 接口一致，即 TracedModule IR 描述的模型是由更加接近用戶視角的高層 op 構(gòu)成，這使得用戶對(duì)模型進(jìn)行一些分析和優(yōu)化時(shí)更加的容易，例如：

• 對(duì) MegEngine 的 python 用戶更加友好，熟悉 MegEngine python 接口的用戶便天然的熟悉了由 TracedModule 表示的模型

• 轉(zhuǎn)換出的模型結(jié)構(gòu)可視化時(shí)更為干凈清晰，用戶很容易的便可將轉(zhuǎn)換后的模型結(jié)構(gòu)與模型源碼對(duì)應(yīng)

• 對(duì)模型進(jìn)行分析，優(yōu)化和向第三方推理框架轉(zhuǎn)換時(shí)更容易，比如：模型量化，算子融合，轉(zhuǎn)換器等

干凈的模型表示

更高層 op 的粒度表示使得模型源碼轉(zhuǎn)換為 TracedModule 后的模型結(jié)構(gòu)更加干凈清晰，用戶很容易的便可以將轉(zhuǎn)換后的模型結(jié)構(gòu)與模型源碼進(jìn)行對(duì)應(yīng)，便于用戶對(duì)模型進(jìn)行分析和調(diào)試。

這里以一個(gè)常用的激活函數(shù) relu6 為例，該激活函數(shù)在 MegEngine 中的 python 接口如下所示：

relu6 函數(shù)在 MegEngine 底層實(shí)際上是調(diào)用了兩個(gè)算子，分別是模式為 MAX 和 MIN 的 Elemwise 算子，熟悉 MegEngine python 源碼的同學(xué)應(yīng)該能夠從上面的代碼中看出 relu6 的前向?qū)崿F(xiàn)里調(diào)用了兩個(gè) elemwise 算子，如圖 2 所示。

如果將一個(gè)調(diào)用了 relu6 函數(shù)導(dǎo)出至由底層 op 所構(gòu)成的模型，其可視化結(jié)果將會(huì)如圖3 所示，可以看到?relu6?變成了兩個(gè) Elmwise 算子，在這個(gè)結(jié)構(gòu)中我們看不到任何關(guān)于?relu6?的信息，不熟悉 MegEngine 底層源碼的用戶面對(duì)這樣一個(gè)模型是比較懵的。

?但如果將該模型代碼轉(zhuǎn)化至 TracedModule 后，將會(huì)得到如圖 4 這樣一個(gè)模型，可以看到 relu6 這個(gè)激活函數(shù)的信息完整的存在于 TracedModule 中，并不會(huì)被轉(zhuǎn)變?yōu)?Elemwise 等其它算子。用戶可以容易的從 TracedModule 中找到與模型源碼所對(duì)應(yīng)的模塊。?

類似 relu6 這樣的 op 在 MegEngine 中還有很多，例如?leaky_relu、interpolate、conv_transpose2d?等都由多個(gè)底層的 op 拼合而成，有些可以從 python 接口的源碼中看出其在底層的實(shí)現(xiàn)，有些卻不太容易看出。可以想象，一個(gè)看起來干凈的使用 MegEngine python 層 op 構(gòu)建的模型代碼，在導(dǎo)出為由框架底層 op 構(gòu)成的模型后，將會(huì)出現(xiàn)模型作者也很難從模型的可視化結(jié)構(gòu)中找到模型某些結(jié)構(gòu)的窘?jīng)r。但將模型源碼導(dǎo)出為由更高層 op 構(gòu)成的 TracedModule 后，將不會(huì)或很少會(huì)出現(xiàn)模型作者不認(rèn)識(shí)可視化出的模型。

直觀的模型圖手術(shù)

將一個(gè) MegEngine 訓(xùn)練出的模型轉(zhuǎn)換至第三方的推理框架進(jìn)行推理時(shí)，常常需要通過圖手術(shù)對(duì)模型結(jié)構(gòu)進(jìn)行一些修改來滿足第三方框架的要求?；?TracedModule 對(duì)模型進(jìn)行修改是非常容易的，如前文提到 TracedModule 模型中的 op 粒度與用戶視角一致，并且構(gòu)成 TracedModule 的基本數(shù)據(jù)結(jié)構(gòu)也都是用戶熟悉的 python 數(shù)據(jù)結(jié)構(gòu)，只需要了解 TracedModule IR 的基本組件，用戶就可以方便的對(duì) TracedModule 所表示的模型運(yùn)行過程進(jìn)行修改。

這里以一個(gè)常用于檢測(cè)模型 Head 模塊中?box?分支的操作為例：

其中?conv?是一個(gè)普通的卷積，scale?和?stride?是兩個(gè)常量 Tensor。在 relu(x)?y 中，當(dāng)y>0 時(shí)，顯然relu(x)?y 與relu(x?y)等價(jià)，所以在轉(zhuǎn)換上述操作時(shí)，常常會(huì)將?scale?和?stride?吸到?conv?的權(quán)重中，吸掉?scale?和?stride?后的模型結(jié)構(gòu)將更加的簡(jiǎn)單，并且也方便轉(zhuǎn)換到一些算子較少的中間模型格式，例如 caffe。在 TracedModule 中我們可以很容易的定位上述操作，并利用?圖手術(shù)接口?完成對(duì)?scale?和?stride?的吸收。圖手術(shù)代碼如下所示。

如圖 5 所示，模型修改完之后，通過?print(graph)?就能直接看到修改之后的圖是否滿足預(yù)期。另外，由于 TracedModule 的 runtime 是 MegEngine 的動(dòng)態(tài)圖，在模型運(yùn)行或圖手術(shù)時(shí)非常的容易調(diào)試。?

寫到這里可能有人會(huì)問，直接修改模型源碼之后再轉(zhuǎn)換豈不是更簡(jiǎn)單？但這會(huì)帶來另外的問題，比如：模型落地過程中，模型可能會(huì)經(jīng)過好幾個(gè)人的處理；引用第三方庫(kù)（例如?basecls?等） 進(jìn)行模型生產(chǎn)時(shí)，直接修改底層源碼顯然是不通用的等。

為了提升用戶體驗(yàn)，TraedMdoule 提供了許多常用的圖手術(shù)接口，并盡可能的使用戶在使用圖手術(shù)接口時(shí)不需要理解和關(guān)注圖內(nèi)部的變化細(xì)節(jié)。在模型圖手術(shù)之后，用戶可以通過打印 Graph 查看圖手術(shù)后的圖是否符合預(yù)期，也可以像運(yùn)行普通 Module 一樣直接運(yùn)行 TracedModule 來查看模型輸出結(jié)果是否正確。能夠使用 MegEngine 構(gòu)建模型的用戶，基本在了解 TracedModule 基本組件后，就可以對(duì) TracedMdoule 模型進(jìn)行圖手術(shù)。另外，我們?yōu)槊恳粋€(gè)圖手術(shù)接口寫了詳細(xì)的使用方法，并提供了一些常見的模型圖手術(shù)?例子?供參考，歡迎大家來試用。

方便的量化模型部署

模型量化是深度學(xué)習(xí)模型部署過程中的一個(gè)重要環(huán)節(jié)，能夠有效減少模型運(yùn)行時(shí)所占用的計(jì)算資源，提高模型的運(yùn)行速度。各個(gè)深度學(xué)習(xí)訓(xùn)練和推理框架都支持對(duì)模型的量化，MegEngine 同樣提供了?模型量化模塊?和豐富的模型量化算法。模型量化的方法大致可以分為以下兩種：

• 量化感知訓(xùn)練（Quantization Aware Training, QAT），一般是在訓(xùn)練時(shí)插入偽量化算子來模擬量化，進(jìn)而緩減量化帶來精度損失

• 訓(xùn)練后量化（Post-Training Quantization, PTQ），一般是利用有限的輸入數(shù)據(jù)對(duì)訓(xùn)練好的模型的權(quán)重和激活值進(jìn)行量化

大多數(shù)的推理框架都支持 PTQ 方法對(duì)模型進(jìn)行量化，用戶只需要提供浮點(diǎn)模型和輸入數(shù)據(jù)集，一般就可以利用框架提供的量化工具完成模型的量化。然而在 PTQ 無法滿足模型的精度的要求時(shí)，便需要借助 MegEngine 等訓(xùn)練框架使用 QAT 方法對(duì)模型進(jìn)行量化，進(jìn)而提高模型量化的精度。

為了更好的支持 MegEngine 量化訓(xùn)練后的模型部署至第三方推理平臺(tái)進(jìn)行推理，MegEngine 團(tuán)隊(duì)開發(fā)了基于 TracedModule 模型轉(zhuǎn)換工具?mgeconvert?來支持量化模型部署到第三方。TracedModule 不僅支持 MegEngine 底層的量化方式，同時(shí)也支持各種自定義的量化算法，這使得基于 TracedModule 導(dǎo)出的量化模型，一般在轉(zhuǎn)換后也能夠滿足目標(biāo)平臺(tái)的量化要求，減小定點(diǎn)模型和偽量化模型之間的差異。用戶只需要將 TracedModule 模型輸入到 mgeconvert 就可以得到以下兩類模型：

• 浮點(diǎn)模型表示（caffe, onnx, tflite）+ 量化參數(shù)文件

• 定點(diǎn)模型表示（tflite）

即 mgeconvert 既支持導(dǎo)出目標(biāo)平臺(tái)的浮點(diǎn)模型和量化參數(shù)文件，也支持導(dǎo)出目標(biāo)平臺(tái)的定點(diǎn)模型。用戶可以方便的使用 MegEengine 量化模塊對(duì)模型進(jìn)行量化，量化后也可以方便的使用 mgeconvert 將模型轉(zhuǎn)到預(yù)期的推理平臺(tái)，mgeconvert 使用方法見?這里，歡迎試用。

總結(jié)

TracedModule 是 MegEngine 設(shè)計(jì)的一種模型格式，設(shè)計(jì)之初便著重考慮了面向用戶視角的 op 粒度，模型圖手術(shù)，量化模型部署等問題，并在文中對(duì)這些問題以及 TracedModule 的效果進(jìn)行了簡(jiǎn)單的介紹。未來 MegEngine 團(tuán)隊(duì)也會(huì)開發(fā)更多基于 TracedModule 的模型發(fā)版工具，例如：模型量化工具，模型優(yōu)化工具等。最后，歡迎大家來試用 TracedModule，也歡迎大家提出建議來一起完善 TracedModule。

附

更多 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ù)證書、定制禮品享不停。