TracedModule: 更友好的模型表示方案橘洞,模型訓(xùn)練到部署的橋梁

作者:曹文剛 | 曠視 MegEngine 架構(gòu)師

TracedModule 介紹

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

圖 1 從一個普通 Module 生成 TracedModule

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)出為 matmuladd 的組合玉转。用戶在使用 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 所示。

圖 2 relu6 的 python 接口和底層實現(xiàn)

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

圖 3 可視化由底層算子構(gòu)成的 relu6

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

圖 4 轉(zhuǎn)化到 TracedModule 中的 relu6

類似 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 是一個普通的卷積画侣,scalestride 是兩個常量 Tensor冰悠。在 relu(x)?y 中,當y>0 時配乱,顯然relu(x)?y 與relu(x?y)等價溉卓,所以在轉(zhuǎn)換上述操作時,常常會將 scalestride 吸到 conv 的權(quán)重中搬泥,吸掉 scalestride 后的模型結(jié)構(gòu)將更加的簡單桑寨,并且也方便轉(zhuǎn)換到一些算子較少的中間模型格式,例如 caffe忿檩。在 TracedModule 中我們可以很容易的定位上述操作尉尾,并利用 圖手術(shù)接口 完成對 scalestride 的吸收。圖手術(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)試。

圖5 圖手術(shù)優(yōu)化前后的 Head 模塊

寫到這里可能有人會問吆豹,直接修改模型源碼之后再轉(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毛甲。

最后編輯于
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
  • 序言:七十年代末玻募,一起剝皮案震驚了整個濱河市,隨后出現(xiàn)的幾起案子跃惫,更是在濱河造成了極大的恐慌,老刑警劉巖蛉顽,帶你破解...
    沈念sama閱讀 222,681評論 6 517
  • 序言:濱河連續(xù)發(fā)生了三起死亡事件先较,死亡現(xiàn)場離奇詭異,居然都是意外死亡噪叙,警方通過查閱死者的電腦和手機霉翔,發(fā)現(xiàn)死者居然都...
    沈念sama閱讀 95,205評論 3 399
  • 文/潘曉璐 我一進店門债朵,熙熙樓的掌柜王于貴愁眉苦臉地迎上來,“玉大人序芦,你說我怎么就攤上這事】矢耍” “怎么了宪塔?”我有些...
    開封第一講書人閱讀 169,421評論 0 362
  • 文/不壞的土叔 我叫張陵,是天一觀的道長比搭。 經(jīng)常有香客問我南誊,道長,這世上最難降的妖魔是什么霉赡? 我笑而不...
    開封第一講書人閱讀 60,114評論 1 300
  • 正文 為了忘掉前任怠苔,我火速辦了婚禮,結(jié)果婚禮上迫肖,老公的妹妹穿的比我還像新娘。我一直安慰自己故爵,他們只是感情好,可當我...
    茶點故事閱讀 69,116評論 6 398
  • 文/花漫 我一把揭開白布诬垂。 她就那樣靜靜地躺著结窘,像睡著了一般充蓝。 火紅的嫁衣襯著肌膚如雪。 梳的紋絲不亂的頭發(fā)上谓苟,一...
    開封第一講書人閱讀 52,713評論 1 312
  • 那天涝焙,我揣著相機與錄音,去河邊找鬼仑撞。 笑死,一個胖子當著我的面吹牛搀缠,可吹牛的內(nèi)容都是我干的近迁。 我是一名探鬼主播簸州,決...
    沈念sama閱讀 41,170評論 3 422
  • 文/蒼蘭香墨 我猛地睜開眼岸浑,長吁一口氣:“原來是場噩夢啊……” “哼!你這毒婦竟也來了璧眠?” 一聲冷哼從身側(cè)響起,我...
    開封第一講書人閱讀 40,116評論 0 277
  • 序言:老撾萬榮一對情侶失蹤袁滥,失蹤者是張志新(化名)和其女友劉穎灾螃,沒想到半個月后,有當?shù)厝嗽跇淞掷锇l(fā)現(xiàn)了一具尸體腰鬼,經(jīng)...
    沈念sama閱讀 46,651評論 1 320
  • 正文 獨居荒郊野嶺守林人離奇死亡熄赡,尸身上長有42處帶血的膿包…… 初始之章·張勛 以下內(nèi)容為張勛視角 年9月15日...
    茶點故事閱讀 38,714評論 3 342
  • 正文 我和宋清朗相戀三年彼硫,在試婚紗的時候發(fā)現(xiàn)自己被綠了。 大學(xué)時的朋友給我發(fā)了我未婚夫和他白月光在一起吃飯的照片溜在。...
    茶點故事閱讀 40,865評論 1 353
  • 序言:一個原本活蹦亂跳的男人離奇死亡他托,死狀恐怖,靈堂內(nèi)的尸體忽然破棺而出志笼,到底是詐尸還是另有隱情,我是刑警寧澤纫溃,帶...
    沈念sama閱讀 36,527評論 5 351
  • 正文 年R本政府宣布紊浩,位于F島的核電站,受9級特大地震影響坊谁,放射性物質(zhì)發(fā)生泄漏滑臊。R本人自食惡果不足惜,卻給世界環(huán)境...
    茶點故事閱讀 42,211評論 3 336
  • 文/蒙蒙 一鬓椭、第九天 我趴在偏房一處隱蔽的房頂上張望颠猴。 院中可真熱鬧翘瓮,春花似錦氧映、人聲如沸。這莊子的主人今日做“春日...
    開封第一講書人閱讀 32,699評論 0 25
  • 文/蒼蘭香墨 我抬頭看了看天上的太陽择份。三九已至烫堤,卻和暖如春,著一層夾襖步出監(jiān)牢的瞬間鸽斟,已是汗流浹背。 一陣腳步聲響...
    開封第一講書人閱讀 33,814評論 1 274
  • 我被黑心中介騙來泰國打工剩燥, 沒想到剛下飛機就差點兒被人妖公主榨干…… 1. 我叫王不留立倍,地道東北人。 一個月前我還...
    沈念sama閱讀 49,299評論 3 379
  • 正文 我出身青樓变擒,卻偏偏與公主長得像寝志,于是被迫代替她去往敵國和親。 傳聞我的和親對象是個殘疾皇子材部,可洞房花燭夜當晚...
    茶點故事閱讀 45,870評論 2 361