地平線軌跡預(yù)測 QCNet 參考算法-V1.0
加入技術(shù)交流群
01 簡介
軌跡預(yù)測任務(wù)的目的是在給定歷史軌跡的情況下預(yù)測未來軌跡。這項(xiàng)任務(wù)在自動(dòng)駕駛、智能監(jiān)控、運(yùn)動(dòng)分析等領(lǐng)域有著廣泛應(yīng)用。傳統(tǒng)方法通常直接利用歷史軌跡來預(yù)測未來，而忽略了預(yù)測目標(biāo)的上下文或查詢信息的影響。這種忽視可能導(dǎo)致預(yù)測精度的下降，特別是在復(fù)雜場景中。
QCNet（Query-Centric Network）引入了一種 query-centric 的預(yù)測機(jī)制，通過對(duì)查詢進(jìn)行顯式建模，增強(qiáng)了對(duì)未來軌跡的預(yù)測能力。首先，通過處理所有場景元素的局部時(shí)空參考框架和學(xué)習(xí)獨(dú)立于全局坐標(biāo)的表示，可以緩存和復(fù)用先前計(jì)算的編碼，另外不變的場景特征可以在所有目標(biāo) agent 之間共享，從而減少推理延遲。其次，使用無錨點(diǎn)查詢來周期性檢測場景上下文，并且在每次重復(fù)時(shí)解碼一小段未來的軌跡點(diǎn)。這種基于查詢的解碼管道將無錨方法的靈活性融入到基于錨點(diǎn)的解決方案中，促進(jìn)了多模態(tài)和長期時(shí)間預(yù)測的準(zhǔn)確性。
本文將介紹軌跡預(yù)測算法 QCNet 在地平線 征程6 平臺(tái)上的優(yōu)化部署。
02 性能精度指標(biāo)
模型參數(shù)：
性能精度表現(xiàn)：
03 公版模型介紹
由于軌跡預(yù)測的歸一化要求，現(xiàn)有方法采用以 agent 為中心的編碼范式來實(shí)現(xiàn)空間旋轉(zhuǎn)平移不變性，其中每個(gè)代理都在由其當(dāng)前時(shí)間步長位置和偏航角確定的局部坐標(biāo)系中編碼。但是觀測窗口每次移動(dòng)時(shí)，場景元素的幾何屬性需要根據(jù) agent 最新狀態(tài)的位置重新歸一化，不斷變化的時(shí)空坐標(biāo)系統(tǒng)阻礙了先前計(jì)算編碼的重用，即使觀測窗口存在很大程度上的重疊。為了解決這個(gè)問題， QCNet 引入了以查詢?yōu)橹行牡木幋a范式，為查詢向量派生的每個(gè)場景元素建立一個(gè)局部時(shí)空坐標(biāo)系，并在其局部參考系中處理查詢?cè)氐奶卣?。然后，在進(jìn)行基于注意力的場景上下文融合時(shí)，將相對(duì)時(shí)空位置注入 Key 和 Value 元素中。下圖展示了場景元素的局部坐標(biāo)系示例：
QCNet 主要由編碼器和****組成，其作用分別為：
編碼器：對(duì)輸入的場景元素進(jìn)行編碼，采用了目前流行的 factorized attention 實(shí)現(xiàn)了時(shí)間維度 attention、Agent-Map cross attention 和 Agent與Agent 間隔的 attention；
****：借鑒 DETR 的****，將編碼器的輸出解碼為每個(gè)目標(biāo) agent 的 K 個(gè)未來軌跡。
3.1 以查詢?yōu)橹行牡膱鼍吧舷挛木幋a
QCNet 首先進(jìn)行了場景元素編碼、相對(duì)位置編碼和地圖編碼，對(duì)于每個(gè) agent 狀態(tài)和 map 上的每個(gè)采樣點(diǎn)，將傅里葉特征與語義屬性(例如：agent 的類別)連接起來，并通過 MLP 進(jìn)行編碼，為了進(jìn)一步生成車道和人行橫道的多邊形級(jí)表示，采用基于注意力的池化對(duì)每個(gè)地圖多邊形內(nèi)采樣點(diǎn)進(jìn)行。這些操作產(chǎn)生形狀為[A, T, D]的 agent 編碼和形狀為[M, D]的 map 編碼，其中 D 表示隱藏的特征維度。為了幫助 agent 編碼捕獲更多信息，編碼器還考慮了跨 agent 時(shí)間 step、agent 之間以及 agent 與 map 之間的注意力并重復(fù)多次。如下圖所示：
3.2 基于查詢的軌跡解碼
軌跡預(yù)測的第二步是利用編碼器輸出的場景編碼來解碼每個(gè)目標(biāo) agent 的 K 個(gè)未來軌跡。受目標(biāo)檢測任務(wù)的啟發(fā)，采用類似 detr 的****來處理這種一對(duì)多問題，并且利用了一個(gè)遞歸的、無錨點(diǎn)的 proposal 模塊來生成自適應(yīng)軌跡錨點(diǎn)，然后是一個(gè)基于錨點(diǎn)的模塊，進(jìn)一步完善初始 proposals。相關(guān)流程如下所示：
04 地平線部署優(yōu)化
整體情況：
QCNet 網(wǎng)絡(luò)主要由 MapEncoder, AgentEncoder, QCDecoder 構(gòu)成，其中 MapEncoder 計(jì)算地圖元素 embedding，AgentEncoder 計(jì)算 agent 元素 embedding，核心組件為 FourierEmbedding 和 AttentionLayer。
改動(dòng)點(diǎn)：
優(yōu)化 FourierEmbedding 結(jié)構(gòu)，去除其中的所有 edge_index，直接計(jì)算形狀為[B, lenq, lenk, D]的相對(duì)信息 r；
將 AttentionLayer 中的 query 形狀設(shè)為[B, lenq, 1, D] , key 形狀為[B, 1, lenk, D], r 形狀為[B, lenq, lenk, D]，利于性能提升;
4.1 性能優(yōu)化
4.1.1 代碼重構(gòu)
FourierEmbedding 將每個(gè)場景元素的極坐標(biāo)轉(zhuǎn)換成傅里葉特征，以方便高頻信號(hào)的學(xué)習(xí)。但是公版 QCNet 使用了大量 edge_index 索引操作， 使得模型中存在大量 BPU 暫不支持的 index_select、scatter 等操作。QCNet 參考算法重構(gòu)了代碼，去除了 FourierEmbedding 中的所有 edge_index，agent_encoder 編碼器注意力層的 query 形狀設(shè)為[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] #重構(gòu)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 #計(jì)算相似性 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)) #重構(gòu)代碼后，edge_index也就不需要了，省去了僅能用CPU運(yùn)行的索引類算子 #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
為了提升性能，主要對(duì) FourierConvEmbedding 做了以下改進(jìn)：
Embedding 和 Linear 層全部替換為了對(duì) BPU 更友好的 Conv1x1；
刪除 self.mlps 層中的 LayerNorm，對(duì)精度基本無影響;
將公版代碼中的torch.stack(continuous_embs).sum(dim=0)直接優(yōu)化為了 add 操作，從而獲得了比較大的性能收益。
對(duì)應(yīng)代碼如下所示：
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 的對(duì)相對(duì)時(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] #取消了公版中對(duì)相對(duì)時(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的輸入中存在距離計(jì)算、torch.norm 等對(duì)量化不友好的操作，為了提升量化精度，將輸入全部置于預(yù)處理中，相關(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 工具獲取量化敏感節(jié)點(diǎn)，然后在 Calibration 和量化訓(xùn)練時(shí)，對(duì) 20% 敏感節(jié)點(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，此操作當(dāng)前僅支持在 CPU 上運(yùn)行，為了提升性能，將其替換為了 torch.where 操作，代碼對(duì)比如下所示：
公版：
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 總結(jié)與建議
5.1 index_select 和 scatter 算子
征程6 僅支持 index_select 和 scatter 索引類算子的 CPU 運(yùn)行，計(jì)算效率較低。QCNet 通過重構(gòu)代碼的形式，優(yōu)化掉了 index_select 和 scatter 操作，實(shí)現(xiàn)了性能的提升。
5.2 ScatterND算子
模型中 nn.embedding 操作引入了目前僅支持在 CPU上 運(yùn)行的 GatherND 算子，后續(xù)將考慮進(jìn)行優(yōu)化。
06 附錄
論文：
公版模型代碼：https://github.com/ZikangZhou/QCNet
*博客內(nèi)容為網(wǎng)友個(gè)人發(fā)布，僅代表博主個(gè)人觀點(diǎn)，如有侵權(quán)請(qǐng)聯(lián)系工作人員刪除。
