本文經(jīng)自動(dòng)駕駛之心公眾號(hào)授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

相信除了少數(shù)自研芯片的大廠,絕大多數(shù)自動(dòng)駕駛公司都會(huì)使用英偉達(dá)NVIDIA芯片,那就離不開(kāi)TensorRT. TensorRT是在NVIDIA各種GPU硬件平臺(tái)下運(yùn)行的一個(gè)C++推理框架。我們利用Pytorch、TF或者其他框架訓(xùn)練好的模型，可以首先轉(zhuǎn)化為onnx格式,再轉(zhuǎn)化為T(mén)ensorRT的格式，然后利用TensorRT推理引擎去運(yùn)行我們這個(gè)模型，從而提升這個(gè)模型在英偉達(dá)GPU上運(yùn)行的速度。

一般來(lái)說(shuō),onnx和TensorRT僅支持相對(duì)比較固定的模型(包括各級(jí)的輸入輸出格式固定,單分支等),最多支持最外層動(dòng)態(tài)輸入(導(dǎo)出onnx可以通過(guò)設(shè)置dynamic_axes參數(shù)確定允許動(dòng)態(tài)變化的維度).但活躍在感知算法前沿的小伙伴們都會(huì)知道,目前一個(gè)重要發(fā)展趨勢(shì)就是端到端(End-2-End),可能涵蓋了目標(biāo)檢測(cè),目標(biāo)跟蹤,軌跡預(yù)測(cè),決策規(guī)劃等全部自動(dòng)駕駛環(huán)節(jié),而且必定是前后幀緊密相關(guān)的時(shí)序模型.實(shí)現(xiàn)了目標(biāo)檢測(cè)和目標(biāo)跟蹤端到端的MUTR3D模型可以作為一個(gè)典型例子(模型介紹可參考:)

實(shí)現(xiàn)真正的端到端多目標(biāo)跟蹤(MOT) --MOTR/MUTR3D中的Label Assignment機(jī)制理論和實(shí)例詳解 https://zhuanlan.zhihu.com/p/609123786

這種模型相當(dāng)于將原來(lái)需要大量后處理和幀間關(guān)聯(lián)的步驟全部放到了模型網(wǎng)絡(luò)里,勢(shì)必帶來(lái)一系列的動(dòng)態(tài)元素,如多if-else分支,子網(wǎng)絡(luò)輸入shape動(dòng)態(tài)變化,和其他一些需要?jiǎng)討B(tài)處理的操作和算子等.這種情況下還能成功轉(zhuǎn)換為T(mén)ensorRT格式并實(shí)現(xiàn)精度對(duì)齊,甚至fp16的精度對(duì)齊嗎?

MUTR3D架構(gòu)因?yàn)檎麄€(gè)過(guò)程涉及多個(gè)細(xì)節(jié),情況各不一樣,縱觀全網(wǎng)的參考資料,甚至google搜索,也很難找到即插即用的方案,只能通過(guò)不斷拆分和實(shí)驗(yàn)來(lái)逐個(gè)解決.通過(guò)博主一個(gè)多月的艱苦探索實(shí)踐(之前對(duì)TensorRT的經(jīng)驗(yàn)不多,沒(méi)有摸清它的脾氣),動(dòng)了不少腦筋,也踩了不少坑,最后終于成功轉(zhuǎn)換并實(shí)現(xiàn)fp32/fp16精度對(duì)齊,且時(shí)延相比單純的目標(biāo)檢測(cè)增加非常小。想在此做一個(gè)簡(jiǎn)單的整理,并為大家提供參考(沒(méi)錯(cuò),一直寫(xiě)綜述,終于寫(xiě)實(shí)踐了!)

1.數(shù)據(jù)格式問(wèn)題

首先是MUTR3D的數(shù)據(jù)格式比較特殊,都是采用實(shí)例形式,這是因?yàn)槊總€(gè)query綁定的信息比較多,都打包成實(shí)例更容易一對(duì)一的存取.但對(duì)于部署而言,輸入輸出只能是tensor,所以首先要對(duì)實(shí)例數(shù)據(jù)進(jìn)行拆解,變成多個(gè)tensor變量.并且由于當(dāng)前幀的query和其他變量是在模型中生成,所以只要輸入前序幀保留的query和其他變量即可,在模型中對(duì)二者進(jìn)行拼接.

2.padding解決輸入動(dòng)態(tài)shape的問(wèn)題

對(duì)于輸入的前序幀query和其他變量,有一個(gè)重要問(wèn)題是shape是不確定的,這是因?yàn)镸UTR3D僅保留前序幀中曾經(jīng)檢出過(guò)目標(biāo)的query.這個(gè)問(wèn)題還是比較容易解決的,最簡(jiǎn)單的辦法就是padding,即padding到一個(gè)固定大小,對(duì)于query可以用全0做padding,數(shù)量具體多少合適,可以根據(jù)自己的數(shù)據(jù)做實(shí)驗(yàn)確定,太少容易漏掉目標(biāo),太多比較浪費(fèi)空間.雖然onnx的dynamic_axes參數(shù)可以實(shí)現(xiàn)動(dòng)態(tài)輸入,但因?yàn)樯婕暗胶罄m(xù)transformer計(jì)算的size,應(yīng)該是有問(wèn)題的,我沒(méi)有嘗試,讀者可以試驗(yàn)一下.

3.padding對(duì)于主transformer中self-attention模塊的影響

如果沒(méi)有使用特殊算子的話,經(jīng)過(guò)padding以后就可以成功轉(zhuǎn)換onnx和tensorrt了.實(shí)際上肯定是有的,但不在本篇的討論范圍,例如MUTR3D中在幀間移動(dòng)reference points時(shí)用到求偽逆矩陣的torch.linalg.inv算子就不支持.如果遇到算子不支持的情況只能先嘗試替換,不行就只能在模型外使用,老司機(jī)的話還可以自己寫(xiě)算子.但因?yàn)檫@一步可以放在模型的預(yù)處理和后處理,我還是選擇把這一步拿到模型外了,自己寫(xiě)算子難度較大.

但是成功轉(zhuǎn)換了就萬(wàn)事大吉了嗎,答案一定是NO,會(huì)發(fā)現(xiàn)精度差距很大.因?yàn)槟Ｐ偷哪K很多,我們先說(shuō)第一個(gè)原因.我們知道在transformer的self-attention階段,會(huì)做多個(gè)query之間的信息交互.而原模型的前序幀只保留了曾經(jīng)檢測(cè)出目標(biāo)的query(模型中稱為active query),應(yīng)該只有這些query與當(dāng)前幀的query進(jìn)行交互.而現(xiàn)在因?yàn)閜adding了很多無(wú)效query,如果所有query一起交互,勢(shì)必會(huì)影響結(jié)果.

解決這個(gè)問(wèn)題受了DN-DETR[1]的啟發(fā),那就是使用attention_mask,在nn.MultiheadAttention中對(duì)應(yīng)'attn_mask'參數(shù),作用就是屏蔽掉不需要進(jìn)行信息交互的query,最初是因?yàn)樵贜LP中每個(gè)句子長(zhǎng)度不一致而設(shè)置的,正好符合我現(xiàn)在的需求,只是需要注意True代表需要屏蔽的query,False代表有效query.

attention mask示意圖 因?yàn)橛?jì)算attention_mask邏輯稍微有點(diǎn)復(fù)雜,很多操作轉(zhuǎn)換TensorRT可能出現(xiàn)新問(wèn)題,所以也應(yīng)該在模型外計(jì)算好之后作為一個(gè)輸入變量輸入模型,再傳遞給transformer.以下是示例代碼:

data['attn_masks'] = attn_masks_init.clone().to(device)
data['attn_masks'][active_prev_num:max_num, :] = True
data['attn_masks'][:, active_prev_num:max_num] = True
[1]DN-DETR: Accelerate DETR Training by Introducing Query DeNoising

4.padding對(duì)于QIM的影響

QIM是MUTR3D中對(duì)transformer輸出的query進(jìn)行的后處理模塊,主要分三步,第一步是篩選active query,即在當(dāng)前幀中檢測(cè)出目標(biāo)的query,依據(jù)是obj_idxs是否>=0(在訓(xùn)練階段還包括隨機(jī)drop query,和隨機(jī)加入fp query,推理階段不涉及),第二步是update query,即針對(duì)第一步中篩選的query做一個(gè)更新,包括query 輸出值的self-attention,ffn,和與query輸入值的shortcut連接,第三步是將更新的query與重新生成的初始query拼接,作為下一幀的輸入.可見(jiàn)第二步中仍然存在我們?cè)诘?點(diǎn)中提到的問(wèn)題,即self-attention不做全部query之間的交互,而是只進(jìn)行active query之間的信息交互.所以在這里又要使用attention mask.

雖然QIM模塊是可選的,但實(shí)驗(yàn)表明對(duì)模型精度的提升是有幫助的.如果要使用QIM的話,這個(gè)attention mask必須在模型里計(jì)算,因?yàn)槟Ｐ屯獠繜o(wú)法得知當(dāng)前幀的檢測(cè)結(jié)果.由于tensorRT的語(yǔ)法限制,很多操作要么會(huì)轉(zhuǎn)換不成功,要么不會(huì)得到想要的結(jié)果,經(jīng)過(guò)多次實(shí)驗(yàn),結(jié)論是直接用索引切片賦值(類似于第3點(diǎn)的示例代碼)操作一般不支持,最好用矩陣計(jì)算的方式,但涉及計(jì)算必須將attention mask的bool類型轉(zhuǎn)為float類型,最后attention mask需要轉(zhuǎn)回bool類型才能使用.以下是實(shí)例代碼:

obj_mask = (obj_idxs >= 0).float()
attn_mask = torch.matmul(obj_mask.unsqueeze(-1), obj_mask.unsqueeze(0)).bool()
attn_mask = ~attn_mask

5.padding對(duì)于輸出結(jié)果的影響

進(jìn)行完以上四點(diǎn),我們基本可以保證模型轉(zhuǎn)換tensorRT的邏輯沒(méi)有問(wèn)題,但輸出結(jié)果經(jīng)過(guò)多次驗(yàn)證后某些幀仍然存在問(wèn)題一度讓我很不解.但一幀幀從數(shù)據(jù)上分析,就會(huì)發(fā)現(xiàn)竟然在某些幀padding的query雖然沒(méi)有參與transformer計(jì)算,卻可以得到一個(gè)較高的score,進(jìn)而得到錯(cuò)誤的結(jié)果.這種情況在數(shù)據(jù)量大的情況下確實(shí)是可能的,因?yàn)閜adding的query只是初始值是0,reference points也是[0,0],與其他隨機(jī)初始化的query進(jìn)行了同樣的操作.但由于畢竟是padding的query,我們并不打算使用他們的結(jié)果,所以必須要進(jìn)行過(guò)濾.

如何過(guò)濾padding query的結(jié)果呢?padding query的標(biāo)志只有他們的索引位置,其他信息都沒(méi)有特異性.而索引信息其實(shí)記錄在第3點(diǎn)使用的attention mask 里,也就是從模型外部傳入的attention mask.這個(gè)mask 是二維的,我們使用其中一維即可(任意一行或任意一列),可以對(duì)padding的track_score直接置為0.記得仍然要注意第4步的注意事項(xiàng),即盡量用矩陣計(jì)算代替索引切片賦值,且計(jì)算必須轉(zhuǎn)換為float類型.代碼示例:

mask = (~attention_mask[-1]).float()
track_scores = track_scores * mask

6.如何動(dòng)態(tài)更新track_id

除了模型主體,其實(shí)還有非常關(guān)鍵的一步,就是動(dòng)態(tài)更新track_id,這也是模型能做到端到端的一個(gè)重要因素.但在原模型中更新track_id的方式是一個(gè)相對(duì)復(fù)雜的循環(huán)判斷, 即高于score thresh且是新目標(biāo)的,賦一個(gè)新的obj_idx, 低于filter score thresh且是老目標(biāo)的,對(duì)應(yīng)的disappear time + 1,如果disappear time超過(guò)miss_tolerance, 對(duì)應(yīng)的obj idx置為-1,即丟棄這個(gè)目標(biāo).

我們知道tensorRT是不支持if-else多分支語(yǔ)句的(好吧,我一開(kāi)始并不知道),這是個(gè)頭疼的問(wèn)題.如果將更新track_id也放到模型外部,不僅影響了模型端到端的架構(gòu),而且也會(huì)導(dǎo)致無(wú)法使用QIM,因?yàn)镼IM篩選query的依據(jù)是更新后的track_id.所以絞盡腦汁也要把更新track_id放到模型里面去.

再次發(fā)揮聰明才智(快用完了),if-else語(yǔ)句也不是不能代替的,比如使用mask并行操作.例如將條件轉(zhuǎn)換為mask(例如tensor[mask] = 0).這里面值得慶幸的是雖然第4,第5點(diǎn)提到tensorRT不支持索引切片賦值操作,但是卻支持bool索引賦值,猜測(cè)可能因?yàn)榍衅僮麟[性改變了tensor的shape吧.但經(jīng)過(guò)多次實(shí)驗(yàn),也不是所有情況下的bool索引賦值都支持的,出現(xiàn)了以下幾種頭疼的情況:

a.賦值的值必須是一個(gè),不能是多個(gè),比如我更新新出現(xiàn)的目標(biāo)時(shí),并不是統(tǒng)一賦值為某一個(gè)id,而是需要為每一個(gè)目標(biāo)賦值連續(xù)遞增的id.這個(gè)想到的辦法是先統(tǒng)一賦值為一個(gè)比較大的不可能出現(xiàn)的數(shù)值,比如1000,避免與之前的id重復(fù),然后在后處理中將1000替換為唯一且連續(xù)遞增的數(shù)值.(我真是個(gè)大聰明)

b.如果要做遞增操作(+=1),只能使用簡(jiǎn)單mask,即不能涉及復(fù)雜邏輯計(jì)算,比如對(duì)disappear_time的更新,本來(lái)需要同時(shí)判斷obj_idx >=0 且 track_scores < 0.35,但由于這兩個(gè)變量都與前面的計(jì)算圖相關(guān),對(duì)他們進(jìn)行與操作可能涉及了比較復(fù)雜的邏輯操作,怎么嘗試都不成功(嘗試了賦值操作代替遞增卻是成功的,無(wú)語(yǔ)).最后的解決辦法是省掉obj_idx >=0 這個(gè)條件.雖然看似不合理,但分析了一下即使將obj_idx=-1的非目標(biāo)的disappear_time遞增,因?yàn)楹罄m(xù)這些目標(biāo)并不會(huì)被選入,所以對(duì)整體邏輯影響不大.

綜上,最后的動(dòng)態(tài)更新track_id示例代碼如下,在后處理環(huán)節(jié)要記得替換obj_idx為1000的數(shù)值.:

def update_trackid(self, track_scores, disappear_time, obj_idxs):
disappear_time[track_scores >= 0.4] = 0
obj_idxs[(obj_idxs == -1) & (track_scores >= 0.4)] = 1000
disappear_time[track_scores < 0.35] += 1
obj_idxs[disappear_time > 5] = -1

至此模型部分的處理就全部結(jié)束了,是不是比較崩潰,但是沒(méi)辦法,部署端到端模型肯定比一般模型要復(fù)雜很多.模型最后會(huì)輸出固定shape的結(jié)果,還需要在后處理階段根據(jù)obj_idx是否>0判斷需要保留到下一幀的query,再根據(jù)track_scores是否>filter score thresh判斷當(dāng)前最終的輸出結(jié)果.總體來(lái)看,需要在模型外進(jìn)行的操作只有三步:幀間移動(dòng)reference_points,對(duì)輸入query進(jìn)行padding,對(duì)輸出結(jié)果進(jìn)行過(guò)濾和轉(zhuǎn)換格式,基本上實(shí)現(xiàn)了端到端的目標(biāo)檢測(cè)+目標(biāo)跟蹤.

還要說(shuō)明的是以上6點(diǎn)存在操作的順序,我這里是按照問(wèn)題分類來(lái)寫(xiě)的,實(shí)際上遇到的順序可能是1->2->3->5->6->4,因?yàn)榈?,6點(diǎn)是使用QIM的前提,第5和第6也存在依賴關(guān)系.還有一個(gè)問(wèn)題是我沒(méi)有使用memory bank,即時(shí)序融合的模塊,因?yàn)榻?jīng)過(guò)實(shí)驗(yàn)這個(gè)模塊提升不是很大,而且對(duì)于端到端跟蹤機(jī)制來(lái)說(shuō),已經(jīng)天然地使用了時(shí)序融合(畢竟直接將前序幀query帶到下一幀),所以時(shí)序融合更加顯得不是非常必要.

好了,現(xiàn)在我們可以進(jìn)行tensorRT的推理結(jié)果和pytorch的推理結(jié)果的對(duì)比,會(huì)發(fā)現(xiàn)fp32精度下可以實(shí)現(xiàn)精度對(duì)齊,撒花!!!!!但如果需要轉(zhuǎn)fp16(可以大幅降低部署時(shí)延),第一次推理會(huì)發(fā)現(xiàn)結(jié)果完全變成none(再次崩潰).導(dǎo)致fp16結(jié)果為none一般都是因?yàn)槌霈F(xiàn)數(shù)據(jù)溢出,即數(shù)值大小超限(fp16最大支持范圍是-65504~+65504),如果你的代碼用了一些自己特殊的操作,或者你的數(shù)據(jù)天然數(shù)值較大,例如內(nèi)外參,pose等數(shù)據(jù)很可能超限,一般通過(guò)縮放等方式解決.這里說(shuō)一下和我以上6點(diǎn)相關(guān)的一個(gè)原因:

7.使用attention_mask導(dǎo)致的fp16結(jié)果為none的問(wèn)題

這個(gè)問(wèn)題非常隱蔽,因?yàn)閱?wèn)題隱藏在torch.nn.MultiheadAttention源碼中,具體在torch.nn.functional.py文件中,有以下幾句:

if attn_mask is not None and attn_mask.dtype == torch.bool:
new_attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
new_attn_mask.masked_fill_(attn_mask, float("-inf"))
attn_mask = new_attn_mask

可以看到,這一步操作是對(duì)attn_mask中值為T(mén)rue的元素用float("-inf")填充,這也是attention mask的原理所在,也就是值為1的位置會(huì)被替換成負(fù)無(wú)窮,這樣在后續(xù)的softmax操作中,這個(gè)位置的輸入會(huì)被加上負(fù)無(wú)窮,輸出的結(jié)果就可以忽略不記,不會(huì)對(duì)其他位置的輸出產(chǎn)生影響.大家也能看出來(lái)了,這個(gè)float("-inf")是fp32精度,肯定超過(guò)fp16支持的范圍了,所以導(dǎo)致結(jié)果為none.我在這里把它替換為fp16支持的下限,即-65504,轉(zhuǎn)fp16就正常了,雖然說(shuō)一般不要修改源碼,但這個(gè)確實(shí)沒(méi)辦法.不要問(wèn)我怎么知道這么隱蔽的問(wèn)題的,因?yàn)椴皇俏乙粋€(gè)人想到的.但如果使用attention_mask之前仔細(xì)研究了原理,想到也不難.

OK,以上就是我踩坑端到端模型部署的全部經(jīng)驗(yàn),說(shuō)全網(wǎng)唯一肯定不是標(biāo)題黨.因?yàn)榻佑|tensorRT也不久,肯定有描述不準(zhǔn)確的地方。

原文鏈接：https://mp.weixin.qq.com/s/EcmNH2to2vXBsdnNvpo0xw