【跟蹤任務 Transformer】Learning Spatio-Temporal Transformer for Visual Tracking
主要思路和創新點
這是一篇為單目標跟蹤任務設計的 Transformer 結構，在各大資料集上都取得了很好的精度，同時速度也大幅超越之前的方法。
其實對於整體網路來說，其中的 Transformer 只是採取了最普遍簡單的形式，甚至不是 Vision Transformer。概括起來其實創新點在於，在網路最後採用了獨特的輸出端融合經過 Transformer 編碼和解碼的序列。之後選擇使用兩階段訓練方法，分別訓練位置預測端和類別預測端。
先單說預測位置的方式，可以看一下整體結構：
搜尋區域和初始模版區域分別輸入共同的骨架網路，這裡使用 ResNet-50 或 ResNet-101，之後將獲得的兩個特徵圖拉長，組合起來輸入 Transformer 模組。可以看下這塊整體結構的程式碼：
def
forward_pass
（
self
，
data
，
run_box_head
，
run_cls_head
）：
feat_dict_list
=
［］
# process the templates
for
i
in
range
（
self
。
settings
。
num_template
）：
template_img_i
=
data
［
‘template_images’
］［
i
］
。
view
（
-
1
，
*
data
［
‘template_images’
］
。
shape
［
2
：］）
# （batch， 3， 128， 128）
template_att_i
=
data
［
‘template_att’
］［
i
］
。
view
（
-
1
，
*
data
［
‘template_att’
］
。
shape
［
2
：］）
# （batch， 128， 128）
feat_dict_list
。
append
（
self
。
net
（
img
=
NestedTensor
（
template_img_i
，
template_att_i
），
mode
=
‘backbone’
））
# process the search regions （t-th frame）
search_img
=
data
［
‘search_images’
］
。
view
（
-
1
，
*
data
［
‘search_images’
］
。
shape
［
2
：］）
# （batch， 3， 320， 320）
search_att
=
data
［
‘search_att’
］
。
view
（
-
1
，
*
data
［
‘search_att’
］
。
shape
［
2
：］）
# （batch， 320， 320）
feat_dict_list
。
append
（
self
。
net
（
img
=
NestedTensor
（
search_img
，
search_att
），
mode
=
‘backbone’
））
# run the transformer and compute losses
seq_dict
=
merge_template_search
（
feat_dict_list
）
#將兩個輸出每個模組分別聯合在一起
out_dict
，
_
，
_
=
self
。
net
（
seq_dict
=
seq_dict
，
mode
=
“transformer”
，
run_box_head
=
run_box_head
，
run_cls_head
=
run_cls_head
）
# out_dict： （B， N， C）， outputs_coord： （1， B， N， C）， target_query： （1， B， N， C）
return
out_dict
def
merge_template_search
（
inp_list
，
return_search
=
False
，
return_template
=
False
）：
‘’‘NOTICE： search region related features must be in the last place’‘’
# 很簡單的操作，搜尋區域在後面因為之後會講到一個提取
seq_dict
=
{
“feat”
：
torch
。
cat
（［
x
［
“feat”
］
for
x
in
inp_list
］，
dim
=
0
），
“mask”
：
torch
。
cat
（［
x
［
“mask”
］
for
x
in
inp_list
］，
dim
=
1
），
“pos”
：
torch
。
cat
（［
x
［
“pos”
］
for
x
in
inp_list
］，
dim
=
0
）}
return
seq_dict
之後保留 Encoder 的輸出及最終 Decoder 輸出，這兩個輸出大小是不一樣的，Decoder 的輸出大小取決於目標序列，大小一般為 1。這兩個輸出通道數是一樣的，因此先將它們點積，再與 Encoder 的輸出做每個畫素的乘法，可以看下面的輸出端圖：
這張圖將輸出端的具體操作解釋的很清楚，在得到乘積後的目標序列後再將其展開為二維的特徵圖。最後將這個特徵圖輸入進兩個結構完全一樣的 FCN 網路，得到兩張熱度圖，分別預測目標物體檢測框的左上角位置和右下角位置。可以參考一下這部分輸出端的程式碼：
def
forward_box_head
（
self
，
hs
，
memory
）：
‘’‘
hs： output embeddings （1， B， N， C） 解碼器的最終輸出
memory： encoder embeddings （H1W1+H2W2， B， C） 編碼器的中間輸出，大小為模板影象素+搜尋區域影象素
’‘’
# adjust shape 這就是為什麼上面在組合時搜尋區域一定要在後面
enc_opt
=
memory
［
-
self
。
feat_len_s
：］
。
transpose
（
0
，
1
）
# encoder output for the search region （B， H2W2， C）
dec_opt
=
hs
。
squeeze
（
0
）
。
transpose
（
1
，
2
）
# （B， C， N）
att
=
torch
。
matmul
（
enc_opt
，
dec_opt
）
# （B， H2W2， N）
opt
=
（
enc_opt
。
unsqueeze
（
-
1
）
*
att
。
unsqueeze
（
-
2
））
。
permute
（（
0
，
3
，
2
，
1
））
。
contiguous
（）
# （B， HW， C， N） ——> （B， N， C， HW）
bs
，
Nq
，
C
，
HW
=
opt
。
size
（）
opt_feat
=
opt
。
view
（
-
1
，
C
，
self
。
feat_sz_s
，
self
。
feat_sz_s
）
# 上面就是輸入 FCN 前的操作
# run the corner head
outputs_coord
=
box_xyxy_to_cxcywh
（
self
。
box_head
（
opt_feat
））
#下面的程式碼會給出
# 上面 xyxy_cxcywh 是從預測兩個頂點轉換成檢測框標註形式
outputs_coord_new
=
outputs_coord
。
view
（
bs
，
Nq
，
4
）
out
=
{
‘pred_boxes’
：
outputs_coord_new
}
return
out
，
outputs_coord_new
這與先前預測檢測框的方法比較有較好的優勢，作者稱是因為能夠顯式化預測中的不確定性，能夠更準確和魯棒。文章給了熱度圖後兩個頂點的計算方式，實際上就是將所有位置與對應值加權平均了：
P 為經過 softmax 的熱度圖；tl 左上角；br 右下角
部分程式碼：
def
conv
（
in_planes
，
out_planes
，
kernel_size
=
3
，
stride
=
1
，
padding
=
1
，
dilation
=
1
）：
# 就是一層 FCN
# 3*3 Conv + BN + ReLU
return
nn
。
Sequential
（
nn
。
Conv2d
（
in_planes
，
out_planes
，
kernel_size
=
kernel_size
，
stride
=
stride
，
padding
=
padding
，
dilation
=
dilation
，
bias
=
True
），
nn
。
BatchNorm2d
（
out_planes
），
nn
。
ReLU
（
inplace
=
True
））
class
Corner_Predictor
（
nn
。
Module
）：
‘’‘產生兩個預測熱度圖及生成頂點’‘’
def
__init__
（
self
，
inplanes
=
64
，
channel
=
256
，
feat_sz
=
20
，
stride
=
16
，
freeze_bn
=
False
）：
super
（
Corner_Predictor
，
self
）
。
__init__
（）
self
。
feat_sz
=
feat_sz
self
。
stride
=
stride
self
。
img_sz
=
self
。
feat_sz
*
self
。
stride
‘’‘top-left corner’‘’
self
。
conv1_tl
=
conv
（
inplanes
，
channel
，
freeze_bn
=
freeze_bn
）
self
。
conv2_tl
=
conv
（
channel
，
channel
//
2
，
freeze_bn
=
freeze_bn
）
self
。
conv3_tl
=
conv
（
channel
//
2
，
channel
//
4
，
freeze_bn
=
freeze_bn
）
self
。
conv4_tl
=
conv
（
channel
//
4
，
channel
//
8
，
freeze_bn
=
freeze_bn
）
self
。
conv5_tl
=
nn
。
Conv2d
（
channel
//
8
，
1
，
kernel_size
=
1
）
‘’‘bottom-right corner’‘’
self
。
conv1_br
=
conv
（
inplanes
，
channel
，
freeze_bn
=
freeze_bn
）
self
。
conv2_br
=
conv
（
channel
，
channel
//
2
，
freeze_bn
=
freeze_bn
）
self
。
conv3_br
=
conv
（
channel
//
2
，
channel
//
4
，
freeze_bn
=
freeze_bn
）
self
。
conv4_br
=
conv
（
channel
//
4
，
channel
//
8
，
freeze_bn
=
freeze_bn
）
self
。
conv5_br
=
nn
。
Conv2d
（
channel
//
8
，
1
，
kernel_size
=
1
）
‘’‘about coordinates and indexs’‘’
with
torch
。
no_grad
（）：
self
。
indice
=
torch
。
arange
（
0
，
self
。
feat_sz
）
。
view
（
-
1
，
1
）
*
self
。
stride
# generate mesh-grid 產生座標
self
。
coord_x
=
self
。
indice
。
repeat
（（
self
。
feat_sz
，
1
））
。
view
（（
self
。
feat_sz
*
self
。
feat_sz
，））
。
float
（）
。
cuda
（）
self
。
coord_y
=
self
。
indice
。
repeat
（（
1
，
self
。
feat_sz
））
。
view
（（
self
。
feat_sz
*
self
。
feat_sz
，））
。
float
（）
。
cuda
（）
def
forward
（
self
，
x
，
return_dist
=
False
，
softmax
=
True
）：
score_map_tl
，
score_map_br
=
self
。
get_score_map
（
x
）
#獲取兩個熱度圖
coorx_tl
，
coory_tl
=
self
。
soft_argmax
（
score_map_tl
）
coorx_br
，
coory_br
=
self
。
soft_argmax
（
score_map_br
）
return
torch
。
stack
（（
coorx_tl
，
coory_tl
，
coorx_br
，
coory_br
），
dim
=
1
）
/
self
。
img_sz
def
get_score_map
（
self
，
x
）：
# top-left branch
x_tl1
=
self
。
conv1_tl
（
x
）
x_tl2
=
self
。
conv2_tl
（
x_tl1
）
x_tl3
=
self
。
conv3_tl
（
x_tl2
）
x_tl4
=
self
。
conv4_tl
（
x_tl3
）
score_map_tl
=
self
。
conv5_tl
（
x_tl4
）
# bottom-right branch
x_br1
=
self
。
conv1_br
（
x
）
x_br2
=
self
。
conv2_br
（
x_br1
）
x_br3
=
self
。
conv3_br
（
x_br2
）
x_br4
=
self
。
conv4_br
（
x_br3
）
score_map_br
=
self
。
conv5_br
（
x_br4
）
return
score_map_tl
，
score_map_br
def
soft_argmax
（
self
，
score_map
，
return_dist
=
False
，
softmax
=
True
）：
‘’‘生成頂點’‘’
score_vec
=
score_map
。
view
（（
-
1
，
self
。
feat_sz
*
self
。
feat_sz
））
# （batch， feat_sz * feat_sz）
prob_vec
=
nn
。
functional
。
softmax
（
score_vec
，
dim
=
1
）
exp_x
=
torch
。
sum
（（
self
。
coord_x
*
prob_vec
），
dim
=
1
）
exp_y
=
torch
。
sum
（（
self
。
coord_y
*
prob_vec
），
dim
=
1
）
return
exp_x
，
exp_y
於是，損失函式就是預測檢測框和標註框的 IoU 損失及 L1 損失：
以上部分均是對跟蹤目標的位置預測，作者提出在注重位置的同時也要注重時間變化對目標特徵帶來的影響。
歸納起來是額外增加了一個訓練階段和訓練分支，這個分支用來評估檢測目標時序上的置信度。該階段為第二階段，鎖定已經訓練好的位置預測主結構，只訓練置信度預測分支。
之後所有的結構會用於驗證和測試階段，輸入變成了三張圖片，額外增加了動態模板。如果在測試過程中預測置信度大於設定閾值，則認為該檢測區域可以作為時間上的動態模板並不斷更新。當然不是每一幀都換，而是隔一定數量比如文章中用了兩百幀。畢竟在長時間序列影片中，較近一幀的特徵相似度是要高一些的，具體結構如下（擴充套件為粉色區域）：
分類埠的結構為三層 MLP，損失函式採用交叉熵：
這部分選取的搜尋區域會有百分之五十的機率為正確圖片，另外會隨機選取同一個影片中的不同幀作為負例，使分類埠判斷是否合理。
實驗結果和視覺化
在 GOT-10k 測試資料集上的實驗結果
在 TrackingNet 測試資料集上的實驗結果
在 VOT2020 資料集上的實驗結果
在 VOT-LT2020 資料集上的實驗結果
在 LaSOT 測試資料集上的精度比較
Transformer Encoder 和 Decoder 的 Attention 結果比較
各結構貢獻的消融實驗比較，實驗中 Position Embedding 採用了 sine
論文資訊
Learning Spatio-Temporal Transformer for Visual Tracking
https：//
arxiv。org/pdf/2103。1715
4。pdf