基于視覺transformer的高效時(shí)空特征學(xué)習(xí)算法

二、背景

高效的時(shí)空建模(Spatiotemporal modeling)是視頻理解和動(dòng)作識別的核心問題。相較于圖像的Transformer網(wǎng)絡(luò)，視頻由于增加了時(shí)間維度，如果將Transformer中的自注意力機(jī)制(Self-Attention)簡單擴(kuò)展到時(shí)空維度，將會導(dǎo)致時(shí)空自注意力高昂的計(jì)算復(fù)雜度和空間復(fù)雜度。許多工作嘗試對時(shí)空自注意力進(jìn)行分解，例如ViViT和Timesformer。這些方法雖然減小了計(jì)算復(fù)雜度，但會引入額外的參數(shù)量。本文提出了一種簡單高效的時(shí)空自注意力Transformer，在對比2D Transformer網(wǎng)絡(luò)不增加計(jì)算量和參數(shù)量情況下，實(shí)現(xiàn)了時(shí)空自注意力機(jī)制。并且在Sthv1&Sthv2, Kinetics400, Diving48取得了很好的性能。

三、方法

視覺Transofrmer通常將圖像分割為不重疊的塊(patch)，patch之間通過自注意力機(jī)制(Self-Attention)進(jìn)行特征聚合，patch內(nèi)部通過全連接層(FFN)進(jìn)行特征映射。每個(gè)Transformer block中，包含Self-Attention和FFN，通過堆疊Transformer block的方式達(dá)到學(xué)習(xí)圖像特征的目的。

在視頻動(dòng)作識別領(lǐng)域，輸入的數(shù)據(jù)是連續(xù)采樣的多幀圖像(常用8幀、16幀、32幀等）學(xué)習(xí)視頻的時(shí)空特征，不僅要學(xué)習(xí)單幀圖像的空間視覺特征，更要建模幀之間的時(shí)域特征。本文提出一種基于視覺transformer的高效時(shí)空特征學(xué)習(xí)算法，具體來說，我們通過將patch按照一定的規(guī)則進(jìn)行移動(dòng)(patch shift)，把當(dāng)前幀中的一部分patch移動(dòng)到其他幀，同時(shí)其他幀也會有一部分patch移動(dòng)到當(dāng)前幀。經(jīng)過patch移動(dòng)之后，對每一幀圖像的patch分別做Self-Attention，這一步學(xué)習(xí)的特征就同時(shí)包含了時(shí)空特征。具體思想可以由下圖所示:

在常用的2D圖像視覺Transformer網(wǎng)絡(luò)結(jié)構(gòu)上，將上述patch shift操作插入到self-attention操作之前即可，無需額外操作，下圖是patch shift transformer block，相比其他視頻transformer的結(jié)構(gòu)，我們的操作不增加額外的計(jì)算量，僅需進(jìn)行內(nèi)存數(shù)據(jù)移動(dòng)操作即可。對于patch shift的移動(dòng)規(guī)則，我們提出幾種設(shè)計(jì)原則:1. 不同幀的塊盡可能均勻地分布。2.合適的時(shí)域感受野。3.保持一定的移動(dòng)塊比例。具體的分析，讀者可以參考正文。

我們對通道移動(dòng)(Channel shift) 與 塊移動(dòng)(patch shift)進(jìn)行了詳盡的分析和討論，這兩種方法的可視化如下:

通道移動(dòng)(Channel shift) 與 塊移動(dòng)(patch shift)都使用了shift操作，但channel shift是通過移動(dòng)所有patch的部分channel的特征來實(shí)現(xiàn)時(shí)域特征的建模，而patch shift是通過移動(dòng)部分patch的全部channel與Self-attention來實(shí)現(xiàn)時(shí)域特征的學(xué)習(xí)?？梢哉J(rèn)為channel shift的時(shí)空建模在空域是稠密的，但在channel上是稀疏的。而patch shift在空域稀疏，在channel上是稠密的。因此兩種方法具有一定的互補(bǔ)性?；诖?，我們提出交替循環(huán)使用 patchshift和channel shift。網(wǎng)絡(luò)結(jié)構(gòu)如下圖所示：

四、實(shí)驗(yàn)結(jié)果

1. 消融實(shí)驗(yàn)

2. 與SOTA方法進(jìn)行對比

3. 運(yùn)行速度

可以看到，PST的實(shí)際推理速度和2D的Swin網(wǎng)絡(luò)接近，但具有時(shí)空建模能力，性能顯著優(yōu)于2D Swin。和Video-Swin網(wǎng)絡(luò)相比，則具有明顯的速度和顯存優(yōu)勢。

4. 可視化結(jié)果

圖中從上到下依次為Kinetics400, Diving48, Sthv1的可視化效果。PST通過學(xué)習(xí)關(guān)聯(lián)區(qū)域的相關(guān)性，并且特征圖能夠反映出視頻當(dāng)中動(dòng)作的軌跡。

審核編輯：郭婷