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

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

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

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

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

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

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

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

2.padding解決輸入動態(tài)shape的問題

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

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

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

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

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

attention mask示意圖 因為計算attention_mask邏輯稍微有點復(fù)雜,很多操作轉(zhuǎn)換TensorRT可能出現(xiàn)新問題,所以也應(yīng)該在模型外計算好之后作為一個輸入變量輸入模型,再傳遞給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對于QIM的影響

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

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

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對于輸出結(jié)果的影響

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

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

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

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

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

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

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

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

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

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

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

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

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

這個問題非常隱蔽,因為問題隱藏在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

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

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

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