Spatial-Temporal Graph Transformer for Multiple Object Tracking

论文地址：Spatial-Temporal Graph Transformer for Multiple Object Tracking

介绍与相关工作

Transformer的成功证明了一种通过自注意力机制建模时间独立性的新范式的出现。最近有工作已经证明了直接使用Transformer建模时空关系的可行性，但是实验表现不是很好。

主要有几个原因：

1. 一个视频包含大量目标，由于没有考虑目标的时空结构，使用一般的Transformer来建模这些目标的时空关系是低效的。
2. 训练一个Transformer去建模长程时空独立性需要很多计算资源以及数据。
3. 基于DETR的检测器对于MOT任务而言并非最优。

基于以上问题，作者提出了STGT网络，通过将追踪目标轨迹视作稀疏赋权图来构建目标间的空间关系。基于此图，STGT构建一个空间图transformer编码器、时间transformer编码器和一个空间transformer解码器。由于稀疏图，所以训练和推理时的计算十分高效。并且由于探索了目标间的结构信息，因此比一般的transformer也更有效。

同时提出了一个级联关联框架来处理低的分的检测和长程遮挡。通过此框架，STGT不需要训练学习相关处理。

早期的工作如Transtrack，使用DETR进行目标检测和特征提取，通过transformer建模轨迹和检测之间的时空关系。作者提出的STGT框架利用空间图transformer建模轨迹和检测之间的时空关系，并且通过分解时空编码器，实现模型有效构建。

概述

主要框架如上，基于DBT范式。框架中包含一个$N_{t-1}$的轨迹集合，每个轨迹代表一个追踪目标。每个轨迹${\mathcal{L}}_i^{t-1}$包含一个状态集合，如其在之前T帧上的位置信息 { x ^ t ′ I } t ′ = t − T t − 1 \{\hat x^I_{t'}\}^{t-1}_{t'=t-T} {x^t′I​}t′=t−Tt−1​和外貌特征 { f ^ t ′ I } t ′ = t − T t − 1 \{\hat f^I_{t'}\}^{t-1}_{t'=t-T} {f^​t′I​}t′=t−Tt−1​。对于新的一帧图片$I_t$，在线追踪算法评估是否追踪目标仍存在或者被遮挡，为已存在轨迹 X ^ t = { x ^ i t } i = 1 N t \hat X_t = \{\hat x^t_{i}\}^{N_t}_{i=1} X^t​={x^it​}i=1Nt​​计算新的位置信息，为新目标产生新轨迹。

如上图，框架主要包括两部分：检测和特征提取子网络，时空图transformer关联子网络。每一帧，检测和特征提取子网络会产生$M_t$个候选目标检测proposal O t = { o j t } j = 1 M t O_t = \{o_j^t\}^{M_t}_{j=1} Ot​={ojt​}j=1Mt​​以及其对应的视觉特征。时空图transformer会为每个轨迹找到最合适的proposal，并且为具体事件建模。

对于每个轨迹$L_i^{t-1}$，最好的匹配是通过计算和proposal $o_t^j$的最大亲密度得分：$\phi (L_i^{t-1},o_j^t)$，$\phi (\cdot )$是一个得分公式，用于计算亲密度。考虑所有的轨迹，则问题可以被公式化为一个优化问题：$$\underset{A^t=\left(a_{ij}^t\right)}{\max }\sum _{i=1}^{N_{t-1}}\sum _{j=1}^{M_t}{a}_{ij}^t\varphi \left({\mathbf{L}}_i^{t-1},o_j^t\right)\text{\hspace{1em}(1)}$$  s.t.  { ∑ i a i j t = 1 , ∀ i = 1 , … , N t − 1 ∑ j a i j t = 1 , ∀ j = 1 , … , M t a i j t ∈ { 0 , 1 } , ∀ i = 1 , … , N t − 1 ; j = 1 , … , M t (2) \text { s.t. }\left\{\begin{array}{ll} \sum_{i} a_{i j}^{t}=1, & \forall i=1, \ldots, N_{t-1} \\ \sum_{j} a_{i j}^{t}=1, & \forall j=1, \ldots, M_{t} \\ a_{i j}^{t} \in\{0,1\}, & \forall i=1, \ldots, N_{t-1} ; j=1, \ldots, M_{t} \end{array}\right. \tag{2}  s.t. ⎩⎨⎧​∑i​aijt​=1,∑j​aijt​=1,aijt​∈{0,1},​∀i=1,…,Nt−1​∀j=1,…,Mt​∀i=1,…,Nt−1​;j=1,…,Mt​​(2)
$A^t=(a_{ij}^t)$表示轨迹 L t − 1 = { L i t − 1 } i = 1 N t L^{t-1} = \{L^{t-1}_i\}^{N_t}_{i=1} Lt−1={Lit−1​}i=1Nt​​和检测候选$O^t$之间的关联。公式2为约束。

公式1和2倍重新整合为一个单独函数：$A^t=\Phi (L^{t-1},O^t)$,$L^{t-1}$和$O^t$分别组成所有的轨迹和候选。

为建模时空关联，作者为每帧的proposals构建了一个赋权空间图${\Theta }^t$，同时还有过去T帧的需追踪目标的赋权空间图集合 Ξ t − 1 = { ξ t − T , ξ 2 , . . . , ξ t − 1 } Ξ^{t−1} = \{ξ^{t−T}, ξ^2, . . . , ξ^{t−1}\} Ξt−1={ξt−T,ξ2,...,ξt−1}。时空图神经网络利用这些图构建一个有效的时空图transformer，来建模追踪目标和新产生的proposal之间的关系。最终会产生一个分配矩阵${\overline{A}}^t$，用于更新追踪目标，同时通过后处理模型对特殊事件（遮挡等）进行处理。

时空图Transformer

时空图Transformer通过图${\Xi }^{t-1}$和${\Theta }^t$学习一个映射$\Phi (\cdot )$建模时空关系，产生分配矩阵${\overline{A}}^t$。该模块包括三部分：空间图Transformer编码层、时间Transformer编码层和一个空间图Transformer解码层。作者通过自注意力机制来提出图多头注意力来建模轨迹和候选之间的空间关系。

时空图Transformer编码器

作者将Transformer分解为两部分，一个为空间图Transformer编码层来建模轨迹间的空间关系，一个为时间Transformer编码层进一步融合和编码轨迹的空间和时间信息。

空间图Transformer编码层

该模块的输入为过去T帧的轨迹状态，用包含轨迹特征的轨迹图序列来表示 Ξ t − 1 = { ξ t − T , ξ 2 , . . . , ξ t − 1 } Ξ^{t−1} = \{ξ^{t−T}, ξ^2, . . . , ξ^{t−1}\} Ξt−1={ξt−T,ξ2,...,ξt−1}，其中 ξ t − 1 = G ( { x i t − 1 } , E X t − 1 , w X t − 1 ) ξ^{t−1} = G(\{x_i^{t-1}\},E_X^{t-1},w_X^{t-1}) ξt−1=G({xit−1​},EXt−1​,wXt−1​)表示t-1帧的轨迹空间图，其中的$x_i^{t-1}$表示该帧第$i$个轨迹的状态，并且如果两个节点相应的边界框的IoU大于0，则存在一条边$E_X^{t-1}$将他们关联，此$w_X^{t-1}$中边的权重设为IoU值。权重矩阵$w_X^{t-1}\in R^{N_{t-1}\times N_{t-1}}$，为一个稀疏矩阵，$(i,j)$表示连接节点$i$和$j$的边的权重，未连接则为0。

轨迹的特征首先通过一个独立于节点的线性层进行嵌入。所有节点特征组成一个特征张量$F_s\in R^{N_{t-1}\times T\times D}$，D表示该线性层的嵌入维度。将其传入时空Transformer编码层如下图。该层内，一个多头图注意力模块用来进行自注意力操作，模块获得特征图$F_s$和图权重$w_X^{t-1}$，生成第$i$个头的自注意力权重：$${\mathcal{F}}_i^{AW}=\mathrm{softmax}\left[\phi \left({\mathcal{F}}_s,W_i^Q,W_i^K\right)\circ w_X^{t-1}\right]\text{\hspace{1em}(3)}$$

$\phi (\cdot )$表示常规缩放的点积，来获得注意力权重，∘表示元素级别的乘法，可以看做是计算每个独立时间戳的空间图的自注意力。

多头图注意力利用图权重$w_X^{t-1}$为有空间交互的轨迹获得非0注意力权重。通过将其注意力集中在一个更小的子集上，空间图transformer编码器层可以更有效地建模交互，并且在训练和推理过程中运行更快

同时对于邻节点使用图网络来获得聚集信息，通过图卷积网络，获得节点特征形成一个value张量$F_i^V$。和公式3中的注意力全中组合，获得最终的多头注意力赋权特征：$${\mathcal{F}}_{att}^{en}=\mathrm{Concate}\left(\left\{{\mathcal{F}}_i^{AW}\otimes {\mathcal{F}}_i^V\right\}\right)\otimes W^O$$
{ ⋅ } \{·\} {⋅}表示迭代和聚合所有注意力头的输出，⊗表示张量mode乘。获得的加权特征通过线性前馈和标准化层获得最终的空间图Transformer编码层的输出。

时间Transformer编码层

轨迹特征会进一步通过一个时间transformer编码层进行编码。会将输出张量的前两个维度进行转变，生成一个结果张量$F_{tmp}^{en}\in R^{T\times N_{t-1}\times D}$。时间transformer编码层通过一个标准的Transformer编码层对每个轨迹的时间维度进行独立编码，沿着时间维度计算自注意力权重，为每个轨迹计算时间注意力赋权特征。

时间transformer编码层的输出就是整个时空图transformer的输出$F_{out}^{en}$。

空间图Transformer解码器

通过获取候选图 Θ t = G ( { o j t } , E O t , w O t ) Θ^t = G(\{o^t_ j\}, E^t_O,w^t_O) Θt=G({ojt​},EOt​,wOt​)和上层输出$F_{out}^{en}$，生成分配矩阵${\overline{A}}^t$。候选图构建和轨迹图相似，每个节点$o_j^t$代表帧$t$的一个候选。对应边界框IoU大于0，则相连接，权重设为IoU。除了真实候选点，还加入虚拟的点，用于对当前帧遮挡或离开的目标负责。实践中，在${\Theta }^t$中加入可学习的embedding$f_{snk}\in R^D$，并且和所有节点以0.5权重相连接。

和之前一样输入图的候选节点的特征需要嵌入和聚集，加入补充节点$f_{snk}$后，或者整体嵌入特征$F_{tgt}^{de}\in R^{(M_t+1)\times 1\times D}$。相似度，送入多头图注意力编码特征，获得最终的带权注意力候选节点特征$F_{att}^{de}\in R^{(M_t+1)\times 1\times D}$。

对于输入的轨迹embedding $F_{out}^{en}$，加入一个虚拟源，处理当前帧$t$中初始化新轨迹的候选，形成一个拓展后的的轨迹embedding $F_{out}^{e{n}^{\prime }}\in R^{T\times (N_{t-1}+1)\times D}$，加入的虚拟源也是个可学习参数。

注意之加入一个虚拟元节点，通过实验证明了一个和多个效果差不多，但是计算效率一个更高。$F_{att}^{de}$复制$N_{t-1}+1$次，得到$F_{att}^{d{e}^{\prime }}\in R^{(M_t+1)\times （N_{t-1}+1）\times D}$。对$F_{att}^{d{e}^{\prime }}$和$F_{out}^{e{n}^{\prime }}$进行多头交叉注意力，生成非标准化注意力权重，输出通过前馈网络和标准化层获得输出张量$R^{(M_t+1)\times （N_{t-1}+1）\times D}$，表示轨迹和候选之间的匹配情况。

空间图解码器的输出可以通过线性层和一个Softmax层获得分配矩阵${\overline{A}}^t\in R^{(M_t+1)\times （N_{t-1}+1）}$。

训练

STGT端到端的训练通过拓展的GT分配矩阵指导。放松公式约束2，实现每个检测必会被和一个轨迹或者虚拟源关联，且一个轨迹可以和多个候选关联：  s.t.  ∑ j N t − 1 + 1 a ˉ i j t = 1 , i ∈ [ 1 , M t ] , a ˉ i j t ∈ { 0 , 1 }  .  \text { s.t. } \quad \sum_{j}^{N_{t-1}+1} \bar{a}_{i j}^{t}=1, i \in\left[1, M_{t}\right], \bar{a}_{i j}^{t} \in\{0,1\} \text { . }  s.t. j∑Nt−1​+1​aˉijt​=1,i∈[1,Mt​],aˉijt​∈{0,1} . 

如此，分配矩阵的一行可以被认为是对于总量为$N_{t-1}+1$分类的分布概率，使用交叉熵优化网络。

在每轮训练迭代过程中，从训练集中随机采样连续的T+1帧序列进行训练，边界框和他们相应ID从帧中获取，通过匹配IoU将GT边界框替换检测器获得的检测框。如此，STGT更加鲁棒，一帧中的所有ID，重新映射为 { 0 , 1 , ⋅ ⋅ ⋅ , N t − 1 } \{0, 1, · · · ,N_{t−1}\} {0,1,⋅⋅⋅,Nt−1​}，显示他们是否匹配轨迹或者虚拟源。

匹配矩阵${\overline{A}}^t$的最后一行代表虚拟点，可能匹配多个轨迹，如此使用一个多标签软边界损失进行独立优化。

总体训练损失定义如下：
$$\begin{array}{rl}\mathcal{L}=& -\frac{1}{M_t}\sum _{m=1}^{{\infty }_t}{y}_m\log \left({\overline{\mathbf{a}}}_m\right)\\ & +\frac{\lambda }{N_{t-1}}\sum _{n=1}^{N_{t-1}}{y}_n^{snk}\log \left(\frac{1}{1+e^{-a_n^{\prime }}}\right)\\ & +\frac{\lambda }{N_{t-1}}\sum _{n=1}^{N_{t-1}}\left(1-y_n^{snk}\right)\log \left(\frac{e^{-a_n^{\prime }}}{1+e^{-a_n^{\prime }}}\right)\end{array}$$
$y_m$和$y_n^{snk}$表示检测候选的ID和虚拟源的ID，${\overline{a}}_m$表示${\overline{A}}^t$的一行， a ‾ M t + 1 = { a n ′ } \overline a_{{M_t}+1} = \{a_n^{'}\} aMt​+1​={an′​}，$\lambda$为权重系数。

级联匹配框架

级联匹配框架如上图：

第一步：使用运动信息匹配和过滤第置信度的候选框，实践中使用卡尔曼滤波预测当前帧中强鲁棒的追踪轨迹，这些轨迹必须成功且连续在过去$K_r$帧匹配。预测的边界框和候选边界框的IoU作为关联得分，使用匈牙利算法进行匹配，IoU高于阈值的匹配对进行关联。剩余的候选置信度低于阈值的直接过滤，剩余的送入STGT进行进一步关联。

第二步：STGT为剩下的轨迹和候选计算拓展分配矩阵${\overline{A}}^t$，其左上角矩阵$\hat{A}\in R^{M_t\times N_{t-1}}$表示真实的轨迹和候选之间的匹配情况，$\hat{A}\in [0,1]$是软分配，通过二分匹配算法产生匹配。相似的这步匹配
时分配得分也需要高于阈值。如此后，未解决的候选框将再次匹配，以恢复遮挡或在第三阶段启动新的目标。

第三步：长程遮挡和重复检测处理模块，由于STGT只建模过去T帧轨迹，过去T帧遮挡轨迹就无法匹配，对于这些轨迹，存储其特征和边界框坐标，当他们可见时，用于计算和候选检测的关联代价。关联代价定义为视觉特征的欧氏距离及被遮挡的轨迹和候选框的归一化顶距离的相加。

$${\mathcal{D}}_{top}=\Vert \left(u_i+\frac{w_i}{2}-u_j-\frac{w_j}{2},v_i-v_j\right)\Vert /h_i$$

$[u,v]$表示边界框左上角的坐标，[w,h]表示边界框大小。

重复检测处理过程中，未关联的候选和所有关联的候选进行匹配，这步的关联代价为关联边界框和未关联边界框在未关联边界框的面积的交集，使用此代价为了删除可能重复进行检测的子边界框，这步的匹配候选检测会被直接去除。

最后，每个保留的候选初始化为新的轨迹，未匹配的轨迹超过$K_p$帧未匹配则直接删除，为超过$K_p$定义为遮挡状态。

实验