作者：曹文剛 | 曠視 MegEngine 架構(gòu)師
TracedModule 介紹
TracedModule 是 MegEngine 中的一種模型格式，用于脫離模型源碼對模型進行訓(xùn)練诡宗、量化锹雏、圖手術(shù)和模型轉(zhuǎn)換，它是模型訓(xùn)練到部署之間的橋梁抠璃。
TracedModule 產(chǎn)生自普通的 Module站楚，它通過以下兩步得到：
- 運行一次 Module，記錄并捕獲模型運行過程中對輸入 Tensor 的所有操作搏嗡，對應(yīng)圖 1 中的
tm.trace_module - 通過一個由 5 種指令所構(gòu)成的 “簡單” 的 high-level IR(intermediate representation) 來描述捕獲到的程序（普通 Module 中的 forward 方法）窿春，對應(yīng)于圖 1 中的
SimpleModule.Graph
TracedModule 的本質(zhì)仍然是一個 Module拉一，它與普通 Module 的區(qū)別在于： 普通 Module 通過用戶實現(xiàn)的 forward 方法描述模型運行過程，而 TracedModule 通過 TracedModule IR 來描述模型的運行過程旧乞。TracedModule IR 由 python 的基本數(shù)據(jù)類型以及 Node 和 Expr 構(gòu)成蔚润，其中：Node 用來表示一個 Tensor 或 Module，并記錄了 Tensor 或 Module 的一些信息尺栖；Expr 用來表示對 Tensor 或 Module 的操作嫡纠，它的輸入和輸出都是 Node。
TracedMdoule IR 中的 Expr 共有以下 5 種：
| OP | 含義 | 例子 |
|---|---|---|
| Input | 表示 Module 的輸入延赌，起到占位的作用 | \ |
| Constant | 表示產(chǎn)生一個常量 Tensor | mge.Tensor([1]) -> %2: const_tensor = Constant() -> (Tensor) |
| GetAttr | 表示獲取 Module 的屬性 | self.linear -> %5: linear = getattr(self, "linear") -> (Linear) |
| CallMethod | 表示調(diào)用 Module 的 forward 方法或 Tensor 的一些方法 | x + self.param -> %7: add_out_1 = relu_out.__add__(param, ) |
| CallFunction | 調(diào)用一個函數(shù) | F.relu(x) -> %4: relu_out = nn.relu(add_out, ) |
通過以上 5 種 Expr 即可表示絕大部分模型的運行過程除盏。
為什么要有 TracedModule？
如前文的介紹挫以，TracedModule IR 是 TracedModule 中的核心數(shù)據(jù)結(jié)構(gòu)者蠕，它用來描述深度學(xué)習(xí)模型的計算過程，使模型能夠脫離源碼而存在掐松。不同的深度學(xué)習(xí)訓(xùn)練框架都有各自的 IR 描述模型踱侣，例如：PyTorch 中的 TorchScript，MindSpore 中的 MindIR甩栈，以及 onnx 等泻仙。這些 IR 大都是一些相對 low-level 的 IR，在模型源碼向這些 IR 轉(zhuǎn)換時常常會發(fā)生 python 的層 op 被轉(zhuǎn)換為多個框架底層 op 組合的現(xiàn)象量没，例如 pytorch 中的 F.Linear 算子在導(dǎo)出到 TorchScript 時可能被導(dǎo)出為 matmul 和 add 的組合玉转。用戶在使用 low-levle 的 IR 表達的模型時會有很多的問題，例如：
- 不了解底層算子用戶可能會很難從模型的可視化結(jié)構(gòu)上與模型源碼對應(yīng)
- 普通用戶學(xué)習(xí) IR 結(jié)構(gòu)較為困難殴蹄，很難對模型進行圖手術(shù)（修改模型圖結(jié)構(gòu)）或優(yōu)化
這種 low-level 的 IR 表達能力往往更加完備究抓，相應(yīng)的也導(dǎo)致 IR 結(jié)構(gòu)極其復(fù)雜，失去高層語義袭灯，使得用戶難以做變換和優(yōu)化刺下，對模型設(shè)計者非常不友好。業(yè)界也提出了一些更加 high-level 的 IR 來解決這些問題稽荧，比如 torch.fx 和 pnnx 等橘茉，這些 IR 都對結(jié)構(gòu)進行了簡化，讓 IR 的描述模型中的 op 粒度更高姨丈，更貼近算法工程師的視角畅卓，使得用戶學(xué)習(xí)更簡單，處理模型也更容易蟋恬。
在 MegEngine 中翁潘，由多個底層 op 組合成的 python 層 op 更多，例如 "resize", "relu6", "softmax" 等歼争，如果直接通過底層 op 表達模型拜马，將會出現(xiàn)導(dǎo)出的模型結(jié)構(gòu)誰也不認識的窘?jīng)r渗勘。為了解決這些問題，MegEngine 參考 torch.fx 和 TorchScript 方案俩莽，改進得到 TracedModule 方案旺坠。TracedModule 的 IR 是一個 high-level 的 IR，它所描述的 op 粒度基本與 MegEngine 的 python 層的 op 一致豹绪，模型中的 op 粒度與用戶視角一致价淌，用戶可以很容易地基于 TracedMdoule IR 對模型進行分析申眼，優(yōu)化和轉(zhuǎn)換瞒津。另外前文提到 TracedMdoule 的本質(zhì)是一個 Module，用戶也可以方便地使用 MegEngine 的模型訓(xùn)練接口對 TracedModule 模型進行訓(xùn)練或參數(shù)微調(diào)括尸。
TracedModule 好在哪巷蚪？
TracedModule 全部由 python 層的數(shù)據(jù)結(jié)構(gòu)構(gòu)成，trace_module 函數(shù)在捕獲用戶代碼的運行邏輯時僅記錄模型中使用的 MegEngine python 層的 function 或 Module濒翻，這使得 TracedModule IR 所描述的 op 粒度基本與 MegEngine 的 python 接口一致屁柏，即 TracedModule IR 描述的模型是由更加接近用戶視角的高層 op 構(gòu)成，這使得用戶對模型進行一些分析和優(yōu)化時更加的容易有送，例如：
- 對 MegEngine 的 python 用戶更加友好淌喻，熟悉 MegEngine python 接口的用戶便天然的熟悉了由 TracedModule 表示的模型
- 轉(zhuǎn)換出的模型結(jié)構(gòu)可視化時更為干凈清晰，用戶很容易的便可將轉(zhuǎn)換后的模型結(jié)構(gòu)與模型源碼對應(yīng)
- 對模型進行分析雀摘，優(yōu)化和向第三方推理框架轉(zhuǎn)換時更容易裸删，比如：模型量化，算子融合阵赠，轉(zhuǎn)換器等
干凈的模型表示
更高層 op 的粒度表示使得模型源碼轉(zhuǎn)換為 TracedModule 后的模型結(jié)構(gòu)更加干凈清晰涯塔，用戶很容易的便可以將轉(zhuǎn)換后的模型結(jié)構(gòu)與模型源碼進行對應(yīng)，便于用戶對模型進行分析和調(diào)試清蚀。
這里以一個常用的激活函數(shù) relu6 為例匕荸，該激活函數(shù)在 MegEngine 中的 python 接口如下所示：
def relu6(x):
relu6 = _get_relu6_op(x.dtype, x.device)
(x,) = relu6(x)
return x
def _get_relu6_op(....)
...
def relu6(inputs, f, c):
(inp,) = inputs[0:1]
max_0 = f("max", inp, c(0))
min_6 = f("min", max_0, c(6))
oup = min_6
(oup_grad,) = yield (oup,)
...
return relu6
...
relu6 函數(shù)在 MegEngine 底層實際上是調(diào)用了兩個算子，分別是模式為 MAX 和 MIN 的 Elemwise 算子枷邪，熟悉 MegEngine python 源碼的同學(xué)應(yīng)該能夠從上面的代碼中看出 relu6 的前向?qū)崿F(xiàn)里調(diào)用了兩個 elemwise 算子榛搔，如圖 2 所示。
如果將一個調(diào)用了 relu6 函數(shù)導(dǎo)出至由底層 op 所構(gòu)成的模型东揣，其可視化結(jié)果將會如圖3 所示践惑，可以看到 relu6 變成了兩個 Elmwise 算子，在這個結(jié)構(gòu)中我們看不到任何關(guān)于 relu6 的信息救斑，不熟悉 MegEngine 底層源碼的用戶面對這樣一個模型是比較懵的童本。
但如果將該模型代碼轉(zhuǎn)化至 TracedModule 后，將會得到如圖 4 這樣一個模型脸候，可以看到 relu6 這個激活函數(shù)的信息完整的存在于 TracedModule 中穷娱，并不會被轉(zhuǎn)變?yōu)?Elemwise 等其它算子绑蔫。用戶可以容易的從 TracedModule 中找到與模型源碼所對應(yīng)的模塊泵额。
類似 relu6 這樣的 op 在 MegEngine 中還有很多，例如 leaky_relu嫁盲、interpolate篓叶、conv_transpose2d 等都由多個底層的 op 拼合而成羞秤，有些可以從 python 接口的源碼中看出其在底層的實現(xiàn)缸托，有些卻不太容易看出●埃可以想象俐镐，一個看起來干凈的使用 MegEngine python 層 op 構(gòu)建的模型代碼，在導(dǎo)出為由框架底層 op 構(gòu)成的模型后哺哼，將會出現(xiàn)模型作者也很難從模型的可視化結(jié)構(gòu)中找到模型某些結(jié)構(gòu)的窘?jīng)r佩抹。但將模型源碼導(dǎo)出為由更高層 op 構(gòu)成的 TracedModule 后，將不會或很少會出現(xiàn)模型作者不認識可視化出的模型取董。
直觀的模型圖手術(shù)
將一個 MegEngine 訓(xùn)練出的模型轉(zhuǎn)換至第三方的推理框架進行推理時棍苹，常常需要通過圖手術(shù)對模型結(jié)構(gòu)進行一些修改來滿足第三方框架的要求∫鹛基于 TracedModule 對模型進行修改是非常容易的枢里，如前文提到 TracedModule 模型中的 op 粒度與用戶視角一致，并且構(gòu)成 TracedModule 的基本數(shù)據(jù)結(jié)構(gòu)也都是用戶熟悉的 python 數(shù)據(jù)結(jié)構(gòu)经窖，只需要了解 TracedModule IR 的基本組件坡垫，用戶就可以方便的對 TracedModule 所表示的模型運行過程進行修改。
這里以一個常用于檢測模型 Head 模塊中 box 分支的操作為例：
F.relu(conv(bbox_subnet) * scale) / stride
其中 conv 是一個普通的卷積画侣，scale 和 stride 是兩個常量 Tensor冰悠。在 relu(x)?y 中，當y>0 時配乱，顯然relu(x)?y 與relu(x?y)等價溉卓，所以在轉(zhuǎn)換上述操作時，常常會將 scale 和 stride 吸到 conv 的權(quán)重中搬泥，吸掉 scale 和 stride 后的模型結(jié)構(gòu)將更加的簡單桑寨，并且也方便轉(zhuǎn)換到一些算子較少的中間模型格式，例如 caffe忿檩。在 TracedModule 中我們可以很容易的定位上述操作尉尾，并利用 圖手術(shù)接口 完成對 scale 和 stride 的吸收。圖手術(shù)代碼如下所示燥透。
graph = traced_head.graph
# 由 conv 的權(quán)重吸收 sacle 和 stride
traced_head.conv.weight *= (traced_head.scale / traced_head.stride)
traced_head.conv.bias *= (traced_head.scale / traced_head.stride)
# 移除 Graph 中的乘 scale
mul_expr = graph.get_expr_by_id(5).as_unique()
graph.replace_node({mul_expr.outputs[0]: mul_expr.inputs[0]})
# 移除 Graph 中的除 stride
div_expr = graph.get_expr_by_id(8).as_unique()
graph.replace_node({div_expr.outputs[0]: div_expr.inputs[0]})
# 刪除 Graph 中無有用的 expr
graph.compile()
如圖 5 所示沙咏，模型修改完之后辨图，通過 print(graph) 就能直接看到修改之后的圖是否滿足預(yù)期。另外肢藐，由于 TracedModule 的 runtime 是 MegEngine 的動態(tài)圖故河，在模型運行或圖手術(shù)時非常的容易調(diào)試。
寫到這里可能有人會問吆豹，直接修改模型源碼之后再轉(zhuǎn)換豈不是更簡單鱼的？但這會帶來另外的問題，比如：模型落地過程中痘煤，模型可能會經(jīng)過好幾個人的處理凑阶；引用第三方庫（例如 basecls 等） 進行模型生產(chǎn)時速勇，直接修改底層源碼顯然是不通用的等坎拐。
為了提升用戶體驗，TraedMdoule 提供了許多常用的圖手術(shù)接口都伪，并盡可能的使用戶在使用圖手術(shù)接口時不需要理解和關(guān)注圖內(nèi)部的變化細節(jié)积担。在模型圖手術(shù)之后，用戶可以通過打印 Graph 查看圖手術(shù)后的圖是否符合預(yù)期先誉，也可以像運行普通 Module 一樣直接運行 TracedModule 來查看模型輸出結(jié)果是否正確的烁。能夠使用 MegEngine 構(gòu)建模型的用戶，基本在了解 TracedModule 基本組件后铃芦，就可以對 TracedMdoule 模型進行圖手術(shù)襟雷。另外，我們?yōu)槊恳粋€圖手術(shù)接口寫了詳細的使用方法咧虎，并提供了一些常見的模型圖手術(shù) 例子 供參考，歡迎大家來試用砰诵。
方便的量化模型部署
模型量化是深度學(xué)習(xí)模型部署過程中的一個重要環(huán)節(jié)，能夠有效減少模型運行時所占用的計算資源胧砰，提高模型的運行速度。各個深度學(xué)習(xí)訓(xùn)練和推理框架都支持對模型的量化偿乖，MegEngine 同樣提供了 模型量化模塊 和豐富的模型量化算法哲嘲。模型量化的方法大致可以分為以下兩種：
- 量化感知訓(xùn)練（Quantization Aware Training, QAT），一般是在訓(xùn)練時插入偽量化算子來模擬量化画切，進而緩減量化帶來精度損失
- 訓(xùn)練后量化（Post-Training Quantization, PTQ）囱怕，一般是利用有限的輸入數(shù)據(jù)對訓(xùn)練好的模型的權(quán)重和激活值進行量化
大多數(shù)的推理框架都支持 PTQ 方法對模型進行量化，用戶只需要提供浮點模型和輸入數(shù)據(jù)集熄捍，一般就可以利用框架提供的量化工具完成模型的量化台丛。然而在 PTQ 無法滿足模型的精度的要求時，便需要借助 MegEngine 等訓(xùn)練框架使用 QAT 方法對模型進行量化防嗡，進而提高模型量化的精度侠坎。
為了更好的支持 MegEngine 量化訓(xùn)練后的模型部署至第三方推理平臺進行推理，MegEngine 團隊開發(fā)了基于 TracedModule 模型轉(zhuǎn)換工具 mgeconvert 來支持量化模型部署到第三方荣德。TracedModule 不僅支持 MegEngine 底層的量化方式童芹，同時也支持各種自定義的量化算法，這使得基于 TracedModule 導(dǎo)出的量化模型署咽，一般在轉(zhuǎn)換后也能夠滿足目標平臺的量化要求，減小定點模型和偽量化模型之間的差異宁否。用戶只需要將 TracedModule 模型輸入到 mgeconvert 就可以得到以下兩類模型：
- 浮點模型表示（caffe, onnx, tflite）+ 量化參數(shù)文件
- 定點模型表示（tflite）
即 mgeconvert 既支持導(dǎo)出目標平臺的浮點模型和量化參數(shù)文件慕匠，也支持導(dǎo)出目標平臺的定點模型。用戶可以方便的使用 MegEengine 量化模塊對模型進行量化台谊，量化后也可以方便的使用 mgeconvert 將模型轉(zhuǎn)到預(yù)期的推理平臺，mgeconvert 使用方法見 這里酪呻，歡迎試用盐须。
總結(jié)
TracedModule 是 MegEngine 設(shè)計的一種模型格式，設(shè)計之初便著重考慮了面向用戶視角的 op 粒度阶冈，模型圖手術(shù)立帖，量化模型部署等問題，并在文中對這些問題以及 TracedModule 的效果進行了簡單的介紹。未來 MegEngine 團隊也會開發(fā)更多基于 TracedModule 的模型發(fā)版工具灌旧，例如：模型量化工具，模型優(yōu)化工具等描融。最后衡蚂，歡迎大家來試用 TracedModule，也歡迎大家提出建議來一起完善 TracedModule毛甲。
