作者:曹文剛 | 曠視 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毛甲。