地平線(xiàn)軌跡預測 QCNet 參考算法-V1.0
加入技術(shù)交流群
掃碼加入
和技術(shù)大咖面對面交流
海量資料庫查詢(xún)
01 簡(jiǎn)介
軌跡預測任務(wù)的目的是在給定歷史軌跡的情況下預測未來(lái)軌跡。這項任務(wù)在自動(dòng)駕駛、智能監控、運動(dòng)分析等領(lǐng)域有著(zhù)廣泛應用。傳統方法通常直接利用歷史軌跡來(lái)預測未來(lái)，而忽略了預測目標的上下文或查詢(xún)信息的影響。這種忽視可能導致預測精度的下降，特別是在復雜場(chǎng)景中。
QCNet（Query-Centric Network）引入了一種 query-centric 的預測機制，通過(guò)對查詢(xún)進(jìn)行顯式建模，增強了對未來(lái)軌跡的預測能力。首先，通過(guò)處理所有場(chǎng)景元素的局部時(shí)空參考框架和學(xué)習獨立于全局坐標的表示，可以緩存和復用先前計算的編碼，另外不變的場(chǎng)景特征可以在所有目標 agent 之間共享，從而減少推理延遲。其次，使用無(wú)錨點(diǎn)查詢(xún)來(lái)周期性檢測場(chǎng)景上下文，并且在每次重復時(shí)解碼一小段未來(lái)的軌跡點(diǎn)。這種基于查詢(xún)的解碼管道將無(wú)錨方法的靈活性融入到基于錨點(diǎn)的解決方案中，促進(jìn)了多模態(tài)和長(cháng)期時(shí)間預測的準確性。
本文將介紹軌跡預測算法 QCNet 在地平線(xiàn) 征程6 平臺上的優(yōu)化部署。
02 性能精度指標
模型參數：
性能精度表現：
03 公版模型介紹
由于軌跡預測的歸一化要求，現有方法采用以 agent 為中心的編碼范式來(lái)實(shí)現空間旋轉平移不變性，其中每個(gè)代理都在由其當前時(shí)間步長(cháng)位置和偏航角確定的局部坐標系中編碼。但是觀(guān)測窗口每次移動(dòng)時(shí)，場(chǎng)景元素的幾何屬性需要根據 agent 最新?tīng)顟B(tài)的位置重新歸一化，不斷變化的時(shí)空坐標系統阻礙了先前計算編碼的重用，即使觀(guān)測窗口存在很大程度上的重疊。為了解決這個(gè)問(wèn)題， QCNet 引入了以查詢(xún)?yōu)橹行牡木幋a范式，為查詢(xún)向量派生的每個(gè)場(chǎng)景元素建立一個(gè)局部時(shí)空坐標系，并在其局部參考系中處理查詢(xún)元素的特征。然后，在進(jìn)行基于注意力的場(chǎng)景上下文融合時(shí)，將相對時(shí)空位置注入 Key 和 Value 元素中。下圖展示了場(chǎng)景元素的局部坐標系示例：
QCNet 主要由編碼器和****組成，其作用分別為：
編碼器：對輸入的場(chǎng)景元素進(jìn)行編碼，采用了目前流行的 factorized attention 實(shí)現了時(shí)間維度 attention、Agent-Map cross attention 和 Agent與Agent 間隔的 attention；
****：借鑒 DETR 的****，將編碼器的輸出解碼為每個(gè)目標 agent 的 K 個(gè)未來(lái)軌跡。
3.1 以查詢(xún)?yōu)橹行牡膱?chǎng)景上下文編碼
QCNet 首先進(jìn)行了場(chǎng)景元素編碼、相對位置編碼和地圖編碼，對于每個(gè) agent 狀態(tài)和 map 上的每個(gè)采樣點(diǎn)，將傅里葉特征與語(yǔ)義屬性(例如：agent 的類(lèi)別)連接起來(lái)，并通過(guò) MLP 進(jìn)行編碼，為了進(jìn)一步生成車(chē)道和人行橫道的多邊形級表示，采用基于注意力的池化對每個(gè)地圖多邊形內采樣點(diǎn)進(jìn)行。這些操作產(chǎn)生形狀為[A, T, D]的 agent 編碼和形狀為[M, D]的 map 編碼，其中 D 表示隱藏的特征維度。為了幫助 agent 編碼捕獲更多信息，編碼器還考慮了跨 agent 時(shí)間 step、agent 之間以及 agent 與 map 之間的注意力并重復多次。如下圖所示：
3.2 基于查詢(xún)的軌跡解碼
軌跡預測的第二步是利用編碼器輸出的場(chǎng)景編碼來(lái)解碼每個(gè)目標 agent 的 K 個(gè)未來(lái)軌跡。受目標檢測任務(wù)的啟發(fā)，采用類(lèi)似 detr 的****來(lái)處理這種一對多問(wèn)題，并且利用了一個(gè)遞歸的、無(wú)錨點(diǎn)的 proposal 模塊來(lái)生成自適應軌跡錨點(diǎn)，然后是一個(gè)基于錨點(diǎn)的模塊，進(jìn)一步完善初始 proposals。相關(guān)流程如下所示：
04 地平線(xiàn)部署優(yōu)化
整體情況：
QCNet 網(wǎng)絡(luò )主要由 MapEncoder, AgentEncoder, QCDecoder 構成，其中 MapEncoder 計算地圖元素 embedding，AgentEncoder 計算 agent 元素 embedding，核心組件為 FourierEmbedding 和 AttentionLayer。
改動(dòng)點(diǎn)：
優(yōu)化 FourierEmbedding 結構，去除其中的所有 edge_index，直接計算形狀為[B, lenq, lenk, D]的相對信息 r；
將 AttentionLayer 中的 query 形狀設為[B, lenq, 1, D] , key 形狀為[B, 1, lenk, D], r 形狀為[B, lenq, lenk, D]，利于性能提升;
4.1 性能優(yōu)化
4.1.1 代碼重構
FourierEmbedding 將每個(gè)場(chǎng)景元素的極坐標轉換成傅里葉特征，以方便高頻信號的學(xué)習。但是公版 QCNet 使用了大量 edge_index 索引操作， 使得模型中存在大量 BPU 暫不支持的 index_select、scatter 等操作。QCNet 參考算法重構了代碼，去除了 FourierEmbedding 中的所有 edge_index，agent_encoder 編碼器注意力層的 query 形狀設為[B, lenq, 1, D] , key 形狀為[B, 1, lenk, D], r 形狀為[B, lenq, lenk, D]，相關(guān)代碼如下所示：
def _attn_block( self, x_src, x_dst, r, mask=None, extra_1dim=False, ): B = x_src.shape[0] if extra_1dim: ... else: if x_src.dim() == 4 and x_dst.dim() == 3: lenq, lenk = x_dst.shape[1], x_src.shape[2] kdim1 = lenq qdim = 1 elif x_src.dim() == 3: kdim1 = qdim = 1 lenq = x_dst.shape[1] lenk = x_src.shape[1] #重構q,k,v,rk,rv的shape q = self.to_q(x_dst).view( B, lenq, qdim, self.num_heads, self.head_dim ) # [B，pl, 1, h, d] k = self.to_k(x_src).view( B, kdim1, lenk, self.num_heads, self.head_dim ) # [B，pl, pt, h, d] v = self.to_v(x_src).view( B, kdim1, lenk, self.num_heads, self.head_dim ) # [B，pl, pt, h, d] if self.has_pos_emb: rk = self.to_k_r(r).view( B, lenq, lenk, self.num_heads, self.head_dim ) rv = self.to_v_r(r).view( B, lenq, lenk, self.num_heads, self.head_dim ) if self.has_pos_emb: k = k + rk v = v + rv #計算相似性 sim = q * k sim = sim.sum(dim=-1) #self.scale = head_dim ** -0.5 sim = sim * self.scale # [B, pl, pt, h] if mask is not None: sim = torch.where( mask.unsqueeze(-1), sim, self.quant(torch.tensor(-100.0).to(mask.device)),) attn = torch.softmax(sim, dim=-2) # [B, pl, pt, h] ... if extra_1dim: inputs = out.view(B, ex_dim, -1, self.num_heads * self.head_dim) else: inputs = out.view(B, -1, self.num_heads * self.head_dim) x = torch.cat([inputs, x_dst], dim=-1) g = torch.sigmoid(self.to_g(x)) #重構代碼后，edge_index也就不需要了，省去了僅能用CPU運行的索引類(lèi)算子 #agg = self.propagate(edge_index=edge_index, x_dst=x_dst, q=q, k=k, v=v, r=r) agg = inputs + g * (self.to_s(x_dst) - inputs) return self.to_out(agg)
代碼路徑：/usr/local/lib/python3.10/dist-packages/hat/models/task_modules/qc
net/rattention.py
4.1.2 FourierConvEmbedding
為了提升性能，主要對 FourierConvEmbedding 做了以下改進(jìn)：
Embedding 和 Linear 層全部替換為了對 BPU 更友好的 Conv1x1；
刪除 self.mlps 層中的 LayerNorm，對精度基本無(wú)影響;
將公版代碼中的torch.stack(continuous_embs).sum(dim=0)直接優(yōu)化為了 add 操作，從而獲得了比較大的性能收益。
對應代碼如下所示：
class FourierConvEmbedding(nn.Module): def __init__( self, input_dim: int, hidden_dim: int, num_freq_bands: int ) -> None: super(FourierConvEmbedding, self).__init__() self.input_dim = input_dim self.hidden_dim = hidden_dim #nn.Embedding替換為了Conv1x1 self.freqs = nn.ModuleList( [ nn.Conv2d(1, num_freq_bands, kernel_size=1, bias=False) for _ in range(input_dim)]) #Linear層替換為了Conv1x1 self.mlps = nn.ModuleList( [nn.Sequential( nn.Conv2d( num_freq_bands * 2 + 1, hidden_dim, kernel_size=1), #刪除LayerNorm #nn.LayerNorm(hidden_dim), nn.ReLU(inplace=True), nn.Conv2d(hidden_dim, hidden_dim, kernel_size=1),) for _ in range(input_dim) ] ) #Linear層替換為了Conv1x1 self.to_out = nn.Sequential( LayerNorm((hidden_dim, 1, 1), dim=1), nn.ReLU(inplace=True), nn.Conv2d(hidden_dim, hidden_dim, 1), ) ... def forward( self, continuous_inputs: Optional[torch.Tensor] = None, categorical_embs: Optional[List[torch.Tensor]] = None, ) -> torch.Tensor: if continuous_inputs is None: ... else: continuous_embs = 0 for i in range(self.input_dim): ... if i == 0: continuous_embs = self.mlps[i](x) else: #將stack+sum的操作替換為add continuous_embs = continuous_embs + self.mlps[i](x) # x = torch.stack(continuous_embs, dim=0).sum(dim=0) x = continuous_embs if categorical_embs is not None: #將stack+sum的操作替換為add # x = x + torch.stack(categorical_embs, dim=0).sum(dim=0) x = x + categorical_embs return self.to_out(x)
代碼路徑：/usr/local/lib/python3.10/dist-packages/hat/models/task_modules/qcnet/fourier_embedding.py
4.1.3 RAttentionLayer
為了提升性能，去除 RAttentionLayer 的對相對時(shí)空編碼 r 的 LayerNorm，相關(guān)代碼如下：
class RAttentionLayer(nn.Module): def __init__( self, ... def forward(self, x, r, mask=None, extra_dim=False): if isinstance(x, torch.Tensor): ... else: x_src, x_dst = x ... x = x[1] #取消了公版中對相對時(shí)空編碼r的LayerNorm #if self.has_pos_emb and r is not None: #r = self.attn_prenorm_r(r) attn = self._attn_block( x_src, x_dst, r, mask=mask, extra_1dim=extra_dim ) # [B, pl, h*d] x = x + self.attn_postnorm(attn) x2 = self.ff_prenorm(x) x = x + self.ff_postnorm(self.ff_mlp(x2)) return x
代碼路徑：/usr/local/lib/python3.10/dist-packages/hat/models/task_modules``/qcnet/rattention.py
4.2 量化精度優(yōu)化
4.2.1 FourierConvEmbedding
QCNetMapEncoder和QCNetAgentEncode的輸入中存在距離計算、torch.norm 等對量化不友好的操作，為了提升量化精度，將輸入全部置于預處理中，相關(guān)代碼如下所示：
class QCNetOEAgentEncoderStream(nn.Module): def __init__( self, ... ) -> None: super().__init__() def build_cur_r_inputs(self, data, cur): pos_pl = data["map_polygon"]["position"] / 10.0 orient_pl = data["map_polygon"]["orientation"] pos_a = data["agent"]["position"][:, :, :cur] / 10.0 # [B, A, HT, 2] head_a = data["agent"]["heading"][:, :, :cur] # [B, A, HT] vel = data["agent"]["velocity"][:, :, :cur, : self.input_dim] / 10.0 ... def build_cur_embs(self, data, cur, map_data, x_a_his, categorical_embs): B, A = data["agent"]["valid_mask"].shape[:2] D = self.hidden_dim ST = self.time_span pl_N = map_data["x_pl"].shape[1] mask_a_cur = data["agent"]["valid_mask"][:, :, cur - 1] ....
代碼路徑：/usr/local/lib/python3.10/dist-packages/hat/models/task_modules/qcnet/agent_st_modeule.py
4.2.2 量化配置首先使用 QAT 的精度 debug 工具獲取量化敏感節點(diǎn)，然后在 Calibration 和量化訓練時(shí)，對 20% 敏感節點(diǎn)配置為 int16 量化，相關(guān)代碼如下：
if os.path.exists(sensitive_path2): sensitive_table1 = torch.load(sensitive_path1) sensitive_table2 = torch.load(sensitive_path2) cali_qconfig_setter = ( sensitive_op_calibration_8bit_weight_16bit_act_qconfig_setter( sensitive_table1, ratio=0.2, ), sensitive_op_calibration_8bit_weight_16bit_act_qconfig_setter( sensitive_table2, ratio=0.2, ), default_calibration_qconfig_setter, ) qat_qconfig_setter = ( sensitive_op_qat_8bit_weight_16bit_fixed_act_qconfig_setter( sensitive_table1, ratio=0.2, ), sensitive_op_qat_8bit_weight_16bit_fixed_act_qconfig_setter( sensitive_table2, ratio=0.2, ), default_qat_fixed_act_qconfig_setter, ) print("Load sensitive table!")4.3 不支持算子替換
4.3.1 cumsum
公版模型的QCNetDecoder中使用了 征程6 暫不支持的 torch.cumsum 算子，參考算法中將其替換為了 Conv1x1，相關(guān)代碼如下：
self.loc_cumsum_conv = nn.Conv2d( self.num_future_steps, self.num_future_steps, kernel_size=1, bias=False, ) self.scale_cumsum_conv = nn.Conv2d( self.num_future_steps, self.num_future_steps, kernel_size=1, bias=False, )
代碼路徑：/usr/local/python3.10/dist-packages/hat/models/models/task_moddules/qcnet/qc_decoder.py
4.3.2 取余操作
公版的 AgentEncoder 使用了處于操作“%”用于 wrap_angle，此操作當前僅支持在 CPU 上運行，為了提升性能，將其替換為了 torch.where 操作，代碼對比如下所示：
公版：
def wrap_angle( angle: torch.Tensor, min_val: float = -math.pi, max_val: float = math.pi) -> torch.Tensor: return min_val + (angle + max_val) % (max_val - min_val)
參考算法：
def wrap_angle( angle: torch.Tensor, min_val: float = -math.pi, max_val: float = math.pi ) -> torch.Tensor: angle = torch.where(angle < min_val, angle + 2 * math.pi, angle) angle = torch.where(angle > max_val, angle - 2 * math.pi, angle) return angle
代碼路徑：/usr/local/python3.10/dist-packages/hat/models/models/task_moddules/qcnet/utils.py
05 總結與建議
5.1 index_select 和 scatter 算子
征程6 僅支持 index_select 和 scatter 索引類(lèi)算子的 CPU 運行，計算效率較低。QCNet 通過(guò)重構代碼的形式，優(yōu)化掉了 index_select 和 scatter 操作，實(shí)現了性能的提升。
5.2 ScatterND算子
模型中 nn.embedding 操作引入了目前僅支持在 CPU上 運行的 GatherND 算子，后續將考慮進(jìn)行優(yōu)化。
06 附錄
論文：
公版模型代碼：https://github.com/ZikangZhou/QCNet
*博客內容為網(wǎng)友個(gè)人發(fā)布，僅代表博主個(gè)人觀(guān)點(diǎn)，如有侵權請聯(lián)系工作人員刪除。
