該示例為參考演算法，僅作為在 征程6 上模型部署的設計參考，非量產演算法。
01 簡介
軌跡預測任務的目的是在給定歷史軌跡的情況下預測未來軌跡。這項任務在自動駕駛、智慧監控、運動分析等領域有著廣泛應用。傳統方法通常直接利用歷史軌跡來預測未來，而忽略了預測目標的上下文或查詢資訊的影響。這種忽視可能導致預測精度的下降，特別是在複雜場景中。
QCNet（Query-Centric Network）引入了一種 query-centric 的預測機制，透過對查詢進行顯式建模，增強了對未來軌跡的預測能力。首先，透過處理所有場景元素的區域性時空參考框架和學習獨立於全域性座標的表示，可以快取和複用先前計算的編碼，另外不變的場景特徵可以在所有目標 agent 之間共享，從而減少推理延遲。其次，使用無錨點查詢來週期性檢測場景上下文，並且在每次重複時解碼一小段未來的軌跡點。這種基於查詢的解碼管道將無錨方法的靈活性融入到基於錨點的解決方案中，促進了多模態和長期時間預測的準確性。
本文將介紹軌跡預測演算法 QCNet 在地平線 征程6 平臺上的最佳化部署。
02 效能精度指標
模型引數：
效能精度表現：
03 公版模型介紹
由於軌跡預測的歸一化要求，現有方法採用以 agent 為中心的編碼正規化來實現空間旋轉平移不變性，其中每個代理都在由其當前時間步長位置和偏航角確定的區域性座標系中編碼。但是觀測視窗每次移動時，場景元素的幾何屬性需要根據 agent 最新狀態的位置重新歸一化，不斷變化的時空座標系統阻礙了先前計算編碼的重用，即使觀測視窗存在很大程度上的重疊。為了解決這個問題， QCNet 引入了以查詢為中心的編碼正規化，為查詢向量派生的每個場景元素建立一個區域性時空座標系，並在其區域性參考系中處理查詢元素的特徵。然後，在進行基於注意力的場景上下文融合時，將相對時空位置注入 Key 和 Value 元素中。下圖展示了場景元素的區域性座標系示例：
QCNet 主要由編碼器和解碼器組成，其作用分別為：
- 編碼器：對輸入的場景元素進行編碼，採用了目前流行的 factorized attention 實現了時間維度 attention、Agent-Map cross attention 和 Agent與Agent 間隔的 attention；
- 解碼器：借鑑 DETR 的解碼器，將編碼器的輸出解碼為每個目標 agent 的 K 個未來軌跡。
3.1 以查詢為中心的場景上下文編碼
QCNet 首先進行了場景元素編碼、相對位置編碼和地圖編碼，對於每個 agent 狀態和 map 上的每個取樣點，將傅立葉特徵與語義屬性(例如：agent 的類別)連線起來，並透過 MLP 進行編碼，為了進一步生成車道和人行橫道的多邊形級表示，採用基於注意力的池化對每個地圖多邊形內取樣點進行。這些操作產生形狀為[A, T, D]的 agent 編碼和形狀為[M, D]的 map 編碼，其中 D 表示隱藏的特徵維度。為了幫助 agent 編碼捕獲更多資訊，編碼器還考慮了跨 agent 時間 step、agent 之間以及 agent 與 map 之間的注意力並重復多次。如下圖所示：
3.2 基於查詢的軌跡解碼
軌跡預測的第二步是利用編碼器輸出的場景編碼來解碼每個目標 agent 的 K 個未來軌跡。受目標檢測任務的啟發，採用類似 detr 的解碼器來處理這種一對多問題，並且利用了一個遞迴的、無錨點的 proposal 模組來生成自適應軌跡錨點，然後是一個基於錨點的模組，進一步完善初始 proposals。相關流程如下所示：
04 地平線部署最佳化
整體情況：
QCNet 網路主要由 MapEncoder, AgentEncoder, QCDecoder 構成，其中 MapEncoder 計算地圖元素 embedding，AgentEncoder 計算 agent 元素 embedding，核心元件為 FourierEmbedding 和 AttentionLayer。
改動點：
- 最佳化 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 效能最佳化
4.1.1 程式碼重構
FourierEmbedding 將每個場景元素的極座標轉換成傅立葉特徵，以方便高頻訊號的學習。但是公版 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]，相關程式碼如下所示：
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執行的索引類運算元
#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/qcnet/rattention.py
4.1.2 FourierConvEmbedding
為了提升效能，主要對 FourierConvEmbedding 做了以下改進：
- Embedding 和 Linear 層全部替換為了對 BPU 更友好的 Conv1x1；
- 刪除 self.mlps 層中的 LayerNorm，對精度基本無影響;
- 將公版程式碼中的
torch.stack(continuous_embs).sum(dim=0)直接最佳化為了 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 的對相對時空編碼 r 的 LayerNorm，相關程式碼如下：
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]
#取消了公版中對相對時空編碼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 量化精度最佳化
4.2.1 FourierConvEmbedding
QCNetMapEncoder 和QCNetAgentEncode 的輸入中存在距離計算、torch.norm 等對量化不友好的操作，為了提升量化精度，將輸入全部置於預處理中，相關程式碼如下所示：
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 工具獲取量化敏感節點，然後在 Calibration 和量化訓練時，對 20% 敏感節點配置為 int16 量化，相關程式碼如下：
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，相關程式碼如下：
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 索引類運算元的 CPU 執行，計算效率較低。QCNet 透過重構程式碼的形式，最佳化掉了 index_select 和 scatter 操作，實現了效能的提升。
5.2 ScatterND運算元
模型中 nn.embedding 操作引入了目前僅支援在 CPU上 執行的 GatherND 運算元，後續將考慮進行最佳化。
06 附錄
- 論文：QCNet
- 公版模型程式碼：https://github.com/ZikangZhou/QCNet