【Paper Note】ViViT: A Video Vision Transformer

Abstract

文章主要transformer在包含时序信息维度的视频格式上的问题展开：

• 视频格式数据生成的token序列数量过多，带来繁重的计算冗余。
• 训练Transfomer结构模型需要引入大规模的数据集，训练对数据条件十分苛刻。

为了高效处理视频数据中生成的大规模时空tokens
①文章提出并探讨了几种对空间和时间维度进行分解的方法，进而提出了相应的网络结构，从而增加模型对视频数据特征提取的效率和可扩展性。
②其次，规范了模型的训练（主要针对模型的训练策略）。目的在小数据集上也能使得Transformer类模型能有很好的效果

Overview of vision transformer 回顾ViT

基础的ViT模型主要有三个模块组成

• Linear Project of Flattened Patches即为Embedding层，对输入的三通道图像数据利用conv卷积层进行分块并完成对应的线性映射，如上式当中的E，而后通过torch.view()进行展平压缩维度。拼接上类别token后采用矩阵相加方式引入位置编码。
• Transformer Encoder模块，对Embedding层输出的token进行多头注意力计算和多层感知机（中间包含Layer Norm）。其中MSA是整个模型的核心部分。
• MLP Head层，堆叠的Transformer Block最终的输出经过Head结构提取出类别token所对应的结果信息，文中通过两个线形层叠加中间插入一个tanh激活函数来实现。

Embedding video clips 视频编码方式

一个视频V有4个维度，T * H * W * C。 变成一个序列token就是 Nt * Nh * Nw * d。加上位置编码， 变成transformer的输入 N * d。

区别于常规的二维图像数据，视频数据相当于需在三维空间内进行采样（拓展了一个时间维度）。而文章中所提出的两钟视频嵌入方法目的都是将视频数据$\mathrm{V}\in {\mathbb{R}}^{\mathrm{T}\times \mathrm{H}\times \mathrm{W}\times \mathrm{C}}$映射到token当中得到$\tilde{\mathrm{z}}\in {\mathbb{R}}^{{\mathrm{n}}_{\mathrm{t}}\times {\mathrm{n}}_{\mathrm{h}}\times {\mathrm{n}}_{\mathrm{w}}\times \mathrm{d}}$，而后添加位置编码并对token进行reshape得到最终Transformer的输入$\mathrm{z}\in {\mathbb{R}}^{\mathrm{N}\times \mathrm{d}}$

Uniform frame sampling 均匀采样

就是先提取帧，然后每一帧按照ViT的方法提取token，然后把不同帧的token拼接起来作为输入

采用相同的采样帧率，从视频数据当中均匀采样$n_t$ 帧，使用相同的embedding方法独立地处理每一个帧当中的patch，而后将得到的所有token拼接concat在一起。具体而言，从每个采样获得的帧当中划分
个不重叠的图像块patch，则共产生$n_w\times n_t$个不重叠的图像块patch，则共产生$n_t\times n_w\times n_t$ 个tokens输入Transformer当中。

然而这种切片方法对于长时间序列的数据来说生成的token长度极大，并且不同帧间首位相连的patch在位置编码上与真实情况不一致。

Tubelet embedding 时空管采样

前一种方法是提取2D图像特征，这种方法是提取立方体，假设每个tublet的shape是t, w, h，那就是说没t帧提取一次特征，取每一帧相同位置的w, hpatch组成输入

从输入volume（体积）当中提取时空上不重叠的“tubes”，这种方法是将vit嵌入到3D的拓展，embedding层就对应的选取三维卷积。则对于维度为$t\times h\times w$的tube管来说，${\mathrm{n}}_{\mathrm{t}}=[\frac{T}{t}],n_{\mathrm{h}}=[\frac{H}{\mathrm{h}}],n_{\mathrm{w}}=[\frac{W}{\mathrm{w}}]$这种采样方法直接在采样的过程当中就融合了时空信息。

提取不重叠，空间-时间的tubes（立方体）。这个tublelt的维度就是： t * h * w。token就包含了时间、宽、高。

所有的模型都是32帧输入的。

看了下vivit_base_k400的config， 模型名：ViViT- B/16*2。其实16 * 16还是ViT一样的方法。

 config.dataset_configs.num_frames = 32  # 采取32帧config.dataset_configs.stride = 2  #2帧为1个config.dataset_configs.crop_size = 224 # 大小224config.model.temporal_encoding_config.method = '3d_conv'
config.model.patches.size = [16, 16, 2]   # H，W是 16* 16的

初始化

模型是以Vit为基础进行训练的，所以初始化需要进行特殊处理

Position emb： 复制t份出来，来适应多帧的处理
Embedding emb：
2d的输入没什么好说的
对于3d的输入，提供了两种不同的方式

下面的公式可以实现在初始的情况下，等价于只用的1帧的情况，参数由模型自己去学习

3D卷积代码介绍

首先2D的卷机是一个平面的卷机（H * W）, 就是一个H * W的平面 和一个 H * W的卷机核，对应点相乘，输出一个值。

那么3D的卷机就是一个立方体（H * W * D）， 就是一个立方体和一个立方体的卷机核相乘，输出一个值。

用pytorch的官方的Conv3D来看， 这个卷机核就是一个3D的立方体 3 * 5 * 2

输入是 (Batch， Channel， Depth， Height， Width) -> （20， 16， 10， 50， 100）

m = nn.Conv3d(16, 33, (3, 5, 2), stride=(2, 1, 1), padding=(4, 2, 0))
input = torch.randn(20, 16, 10, 50, 100)
output = m(input)
output.shape # torch.Size([20, 33, 8, 50, 99]

更改为paper中的输入：

输入是一个batch:16, 3 * 224 * 224的图片， 一共有32帧，

使用kenel，2 * 16 * 16， 理解为2帧变1帧， 图像上 16 * 16的不重叠区域

m = nn.Conv3d(3, 1, (2, 16, 16), stride=(2, 16, 16))
input = torch.randn(16, 3, 32, 224, 224)
output = m(input)
output.shape # torch.Size([16, 1, 16, 14, 14])

视频编码输入到模型当中

输入视频，均匀采样， 知道采样的帧数（n_sampled_frames）， 去算间隔， 采样，输出。

def sample_frames_uniformly(x: jnp.ndarray,n_sampled_frames: int) -> jnp.ndarray:"""Sample frames from the input video."""if x.ndim != 5:raise ValueError('Input shape should be [bs, t, h, w, c].')num_frames = x.shape[1]if n_sampled_frames < num_frames:t_start_idx = num_frames / (n_sampled_frames + 1)t_step = t_start_idxelse:t_start_idx = 0t_step = 1t_end_idx = num_framestemporal_indices = jnp.arange(t_start_idx, t_end_idx, t_step)temporal_indices = jnp.round(temporal_indices).astype(jnp.int32)temporal_indices = jnp.minimum(temporal_indices, num_frames - 1)return x[:, temporal_indices]  # [n, t_s, in_h, in_w, c]

编码后，从batch， time，h, w, c -> batch, thw, c

def temporal_encode(x,temporal_encoding_config,patches,hidden_size,return_1d=True,name='embedding'):"""Encode video for feeding into ViT."""n, _, in_h, in_w, c = x.shapeif temporal_encoding_config.method == 'temporal_sampling':n_sampled_frames = temporal_encoding_config.n_sampled_framesx = video_utils.sample_frames_uniformly(x, n_sampled_frames)t_s = x.shape[1]x = jnp.reshape(x, [n, t_s * in_h, in_w, c])x = embed_2d_patch(x, patches, hidden_size)temporal_dims = t_sif return_1d:n, th, w, c = x.shapex = jnp.reshape(x, [n, th * w, c])else:n, th, w, c = x.shapex = jnp.reshape(x, [n, t_s, -1, w, c])elif temporal_encoding_config.method == '3d_conv':kernel_init_method = temporal_encoding_config.get('kernel_init_method',None)x = embed_3d_patch(x, patches, hidden_size, kernel_init_method, name)temporal_dims = x.shape[1]if return_1d:n, t, h, w, c = x.shapex = jnp.reshape(x, [n, t * h * w, c])else:raise AssertionError('Unknown temporal encoding method.')assert x.size > 0, ('Found zero tokens after temporal encoding. ''Perhaps one of the patch sizes is such that ''floor(dim_size / patch_size) = 0?')return x, temporal_dims

Transformer Models for Video

Spatio-temporal attention 空间-时间注意力

伴随着采样的输入帧数增加，token的数量也会线性增加。运算量会平方倍的增加，所以需要更加有效的结构。

这种模型简单地将所有的tokens(包括时空)简单地通过Transformer encoder层，导致问题就是引入指数增长的计算量，每个Transformer层对所有时空token均进行成对交互，这种方式极其低效，具体模型构成方式与另一篇文章：Video Transformer Net 所提出的结构类似，如图4所示

Factorised encoder

使用两个 transformer

• 第一个是 spatial transformer，输入是某一帧的多个token，输出一个token
• 第二个是temporal transformer，输入是前一步多帧的token（每帧对应一个token），输出结果就通过mlp进行分类

模型是2个单独的transformer encoder组成的：

• 空间编码器，通过对同一时间索引的token建模。输出cls_token。
• 这个帧维度的表征，连接在一起，输入时间编码器中。这个输出就是最后的结果。

Factorised encoder方法：构建两个单独的transformer encoder，分别针对空间和时间处理。首先利用空间编码器（Space Transformer），通过对同一时间索引的token建模。输出cls_token。而后将**输出的类别token和帧维度的表征token拼接输入到时间编码器（Time Transformer）**中得到最终的结果，模型结构如图5所示（相当于两个Transformer模型的叠加），实现代码如下：

class ViViT(nn.Module):def __init__(self, image_size, patch_size, num_classes, num_frames, dim = 192, depth = 4, heads = 3, pool = 'cls', in_channels = 3, dim_head = 64, dropout = 0.,emb_dropout = 0., scale_dim = 4, ):super().__init__()assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'num_patches = (image_size // patch_size) ** 2patch_dim = in_channels * patch_size ** 2self.to_patch_embedding = nn.Sequential(Rearrange('b t c (h p1) (w p2) -> b t (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),nn.Linear(patch_dim, dim),)self.pos_embedding = nn.Parameter(torch.randn(1, num_frames, num_patches + 1, dim))self.space_token = nn.Parameter(torch.randn(1, 1, dim))self.space_transformer = Transformer(dim, depth, heads, dim_head, dim*scale_dim, dropout)self.temporal_token = nn.Parameter(torch.randn(1, 1, dim))self.temporal_transformer = Transformer(dim, depth, heads, dim_head, dim*scale_dim, dropout)self.dropout = nn.Dropout(emb_dropout)self.pool = poolself.mlp_head = nn.Sequential(nn.LayerNorm(dim),nn.Linear(dim, num_classes))def forward(self, x):x = self.to_patch_embedding(x)b, t, n, _ = x.shapecls_space_tokens = repeat(self.space_token, '() n d -> b t n d', b = b, t=t)x = torch.cat((cls_space_tokens, x), dim=2)x += self.pos_embedding[:, :, :(n + 1)]x = self.dropout(x)x = rearrange(x, 'b t n d -> (b t) n d')x = self.space_transformer(x)x = rearrange(x[:, 0], '(b t) ... -> b t ...', b=b)cls_temporal_tokens = repeat(self.temporal_token, '() n d -> b n d', b=b)x = torch.cat((cls_temporal_tokens, x), dim=1)x = self.temporal_transformer(x)x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]return self.mlp_head(x)

Factorised self-attention

通过 self-attention 层将时空数据分开处理

• 空间层只在同一帧内不同token间进行attention操作
• 时间层对不同帧同一位置的token进行attention操作
• 先计算空间自注意力（token中有相同的时间索引），再计算时间的自注意力（token中有相同的空间索引），其实先后顺序无所谓，只要串行就行
  

相较于Model 1，这个模型包含相同数量的Transformer层。而此模型思路不是在第$l$ 层计算所有成对的token $z^l$ 的多头自注意力，而是将自注意力计算在空间和时间上分解，首先只计算空间上的自注意力（对于相同时间索引的token），而后再进行时间上的计算（对于相同的空间索引）。在降低计算复杂度的同时在每个Transformer层均完成了时空层面的自注意力交互。其结构如图6所示。

自注意力计算：

先时间后空间，或者先空间后时间没有区别

def _reshape_to_time_space(x, temporal_dims):if x.ndim == 3:b, thw, d = x.shapeassert thw % temporal_dims == 0hw = thw // temporal_dimsx = jnp.reshape(x, [b, temporal_dims, hw, d])assert x.ndim == 4return x、

reshape_to_2d_factorized, 就是将batch， h * w, channel -> batch, w, h, channel

def reshape_to_2d_factorized(x: jnp.ndarray, axis: int,two_d_shape: Tuple[int, int, int, int]):"""Converts 1d inputs back to 2d after axial attention."""assert x.ndim == 3, ('The input dimention should be ''[batch_size, height*width, channel]')batch_size, height, width, channel = two_d_shapeif axis == 1:assert x.shape[0] == batch_size * widthreturn x.reshape((batch_size, width, height, channel)).transpose((0, 2, 1, 3))elif axis == 2:assert x.shape[0] == batch_size * heightreturn x.reshape(two_d_shapedef reshape_to_2d_factorized(x: jnp.ndarray, axis: int,two_d_shape: Tuple[int, int, int, int]):"""Converts 1d inputs back to 2d after axial attention."""assert x.ndim == 3, ('The input dimention should be ''[batch_size, height*width, channel]')batch_size, height, width, channel = two_d_shapeif axis == 1:assert x.shape[0] == batch_size * widthreturn x.reshape((batch_size, width, height, channel)).transpose((0, 2, 1, 3))elif axis == 2:assert x.shape[0] == batch_size * heightreturn x.reshape(two_d_shape)

在不同的维度上做注意力，来实现时间和空间。

其实也是一样的，LN + atttion + 残差连

def _run_attention_on_axis(inputs, axis, two_d_shape):"""Reshapes the input and run attention on the given axis."""inputs = model_utils.reshape_to_1d_factorized(inputs, axis=axis)x = nn.LayerNorm(dtype=self.dtype, name='LayerNorm_{}'.format(_AXIS_TO_NAME[axis]))(inputs)x = self_attention(name='MultiHeadDotProductAttention_{}'.format(_AXIS_TO_NAME[axis]))(x, deterministic=deterministic)x = nn.Dropout(rate=self.dropout_rate)(x, deterministic)x = x + inputsreturn model_utils.reshape_to_2d_factorized(x, axis=axis, two_d_shape=two_d_shape)

整个这个因式分解的注意力模块，就是在不同的轴上，做自注意力。

可以使用先时间后空间，attention_axes = (1, 2)。或者先空间后时间，attention_axes= (2, 1)。

所以整个就是： 时间attn + 空间attn + LN + MLP

Factorised dot-product attention

时间、空间heads是并行的，而不是串行的。
spatial还是同一帧内不同token，temporal是不同帧同一位置的token

第四种模型的思想则是通过利用dot-product点积注意力操作来取代上述的因式分解factorisation操作，通过注意力计算的方式来代替简单的张量reshape。思想是对于空间注意力和时间注意力分别构建对应的键、值。具体思路如图所示。

class FDATransformerEncoder(nn.Module):"""Factorized Dot-product Attention Transformer Encoder"""def __init__(self, dim, depth, heads, dim_head, mlp_dim, nt, nh, nw, dropout=0.):super().__init__()self.layers = nn.ModuleList([])self.nt = ntself.nh = nhself.nw = nwfor _ in range(depth):self.layers.append(PreNorm(dim, FDAttention(dim, nt, nh, nw, heads=heads, dim_head=dim_head, dropout=dropout)))def forward(self, x):for attn in self.layers:x = attn(x) + xreturn xclass FDAttention(nn.Module):"""Factorized Dot-product Attention"""def __init__(self, dim, nt, nh, nw, heads=8, dim_head=64, dropout=0.):super().__init__()inner_dim = dim_head * headsproject_out = not (heads == 1 and dim_head == dim)self.nt = ntself.nh = nhself.nw = nwself.heads = headsself.scale = dim_head ** -0.5self.attend = nn.Softmax(dim=-1)self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)self.to_out = nn.Sequential(nn.Linear(inner_dim, dim),nn.Dropout(dropout)) if project_out else nn.Identity()def forward(self, x):b, n, d, h = *x.shape, self.headsqkv = self.to_qkv(x).chunk(3, dim=-1)q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), qkv)qs, qt = q.chunk(2, dim=1)ks, kt = k.chunk(2, dim=1)vs, vt = v.chunk(2, dim=1)# Attention over spatial dimensionqs = qs.view(b, h // 2, self.nt, self.nh * self.nw, -1)ks, vs = ks.view(b, h // 2, self.nt, self.nh * self.nw, -1), vs.view(b, h // 2, self.nt, self.nh * self.nw, -1)spatial_dots = einsum('b h t i d, b h t j d -> b h t i j', qs, ks) * self.scalesp_attn = self.attend(spatial_dots)spatial_out = einsum('b h t i j, b h t j d -> b h t i d', sp_attn, vs)# Attention over temporal dimensionqt = qt.view(b, h // 2, self.nh * self.nw, self.nt, -1)kt, vt = kt.view(b, h // 2, self.nh * self.nw, self.nt, -1), vt.view(b, h // 2, self.nh * self.nw, self.nt, -1)temporal_dots = einsum('b h s i d, b h s j d -> b h s i j', qt, kt) * self.scaletemporal_attn = self.attend(temporal_dots)temporal_out = einsum('b h s i j, b h s j d -> b h s i d', temporal_attn, vt)

消融实验

比较不同的token获取方式
比较了不同变种的transformer
比较了多种数据增强方式
比较了不同输入数据尺寸
比较了几类变种
比较了不同的输入帧数