本文經自動駕駛之心公眾號授權轉載，轉載請聯系出處。

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

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

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

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

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

1.數據格式問題

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

2.padding解決輸入動態shape的問題

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

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

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

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

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

attention mask示意圖 因為計算attention_mask邏輯稍微有點復雜,很多操作轉換TensorRT可能出現新問題,所以也應該在模型外計算好之后作為一個輸入變量輸入模型,再傳遞給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進行的后處理模塊,主要分三步,第一步是篩選active query,即在當前幀中檢測出目標的query,依據是obj_idxs是否>=0(在訓練階段還包括隨機drop query,和隨機加入fp query,推理階段不涉及),第二步是update query,即針對第一步中篩選的query做一個更新,包括query 輸出值的self-attention,ffn,和與query輸入值的shortcut連接,第三步是將更新的query與重新生成的初始query拼接,作為下一幀的輸入.可見第二步中仍然存在我們在第3點中提到的問題,即self-attention不做全部query之間的交互,而是只進行active query之間的信息交互.所以在這里又要使用attention mask.

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

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

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

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

6.如何動態更新track_id

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

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

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

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

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

綜上,最后的動態更新track_id示例代碼如下,在后處理環節要記得替換obj_idx為1000的數值.:

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

至此模型部分的處理就全部結束了,是不是比較崩潰,但是沒辦法,部署端到端模型肯定比一般模型要復雜很多.模型最后會輸出固定shape的結果,還需要在后處理階段根據obj_idx是否>0判斷需要保留到下一幀的query,再根據track_scores是否>filter score thresh判斷當前最終的輸出結果.總體來看,需要在模型外進行的操作只有三步:幀間移動reference_points,對輸入query進行padding,對輸出結果進行過濾和轉換格式,基本上實現了端到端的目標檢測+目標跟蹤.

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

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

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

OK,以上就是我踩坑端到端模型部署的全部經驗,說全網唯一肯定不是標題黨.因為接觸tensorRT也不久,肯定有描述不準確的地方。

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