TSM: Temporal Shift Module for Efficient Video Understanding

发表于 更新于 分类于
本文字数： 3.3k 阅读时长 ≈ 6 分钟

原文地址：TSM: Temporal Shift Module for Efficient Video Understanding

官方实现：mit-han-lab/temporal-shift-module

复现地址：ZJCV/TSM

摘要

The explosive growth in video streaming gives rise to challenges on performing video understanding at high accuracy and low computation cost. Conventional 2D CNNs are computationally cheap but cannot capture temporal relationships; 3D CNN based methods can achieve good performance but are computationally intensive, making it expensive to deploy. In this paper, we propose a generic and effective Temporal Shift Module (TSM) that enjoys both high efficiency and high performance. Specifically, it can achieve the performance of 3D CNN but maintain 2D CNN's complexity. TSM shifts part of the channels along the temporal dimension; thus facilitate information exchanged among neighboring frames. It can be inserted into 2D CNNs to achieve temporal modeling at zero computation and zero parameters. We also extended TSM to online setting, which enables real-time low-latency online video recognition and video object detection. TSM is accurate and efficient: it ranks the first place on the Something-Something leaderboard upon publication; on Jetson Nano and Galaxy Note8, it achieves a low latency of 13ms and 35ms for online video recognition. The code is available at: this https URL.

视频流的爆炸性增长给以高精度和低计算成本执行视频理解带来了挑战。传统的2D CNN计算成本低，但无法捕捉时间关系；基于3D CNN的方法可以实现良好的性能，但计算量大，部署成本高。在这篇文章中，我们提出了一个通用有效的时间移位模块，该模块同时具有高效率和高性能。具体来说，它可以达到3D CNN的性能，但保持2D CNN的复杂性。TSM沿时间维度移动部分通道；从而促进相邻帧之间信息交换。它可以插入到2D神经网络中，实现零计算、零参数的时态建模。我们还将TSM扩展到在线设置，支持实时低延迟在线视频识别和视频对象检测。TSM精准高效:一经发布，在Something-Something排行榜上排名第一；在Jetson Nano和Galaxy Note8上，它实现了13毫秒和35毫秒的低延迟，用于在线视频识别。代码可在以下网址获得：https://github.com/mit-han-lab/temporal-shift-module。

引言

假定激活大小\(A\in R^{N\times C\times T\times H\times W}\)，其中\(N\)表示批量大小，\(C\)表示通道数，\(T\)表示时间维度，\(H\)表示高，\(W\)表示宽

  • 对于2D卷积而言，其操作独立于时间维度，如上图（a）所示；
  • TSM（temporal shift module, 时间移位模块）沿着时间维度对通道进行向上/向下移位（移动一位）操作，如上图（b）所示；
  • 对于实时视频推理，无法对未来时间维度进行移位，所以仅对过去时间进行单向移位，如上图（c）所示。

通过移位操作，相邻帧信息与当前帧信息混合在一起，经过2D卷积计算就能够进行时间建模。移位操作的执行会带来两个问题：

  1. 移位操作会发生数据移动，实际计算过程中会造成很大的内存消耗；
  2. 移动太多的通道数会损害空间建模能力，导致模型性能退化。

所以最终的TSM模块实现如下：

  1. 沿着时间维度移动一小部分的通道；
  2. 将TSM插入到残差分支上，保证当前帧的空间特征不会被损害

灵感

以一维卷积为例，假定卷积权重向量\(W=(w_{1}, w_{2}, w_{3})\)，输入\(X\)是一个1D向量，卷积操作\(Y=Conv(W, X)\)实现如下：

\[ Y_{i}=w_{1}X_{i-1} + w_{2}X_{i} + w_{3}X_{i+1} \]

卷积运算可以分两步进行：移位操作和乘-加操作。首先将卷积核中心移动到输入\(X\)的第\(i\)个位置：

\[ X_{i}^{-1} = X_{i-1}, X_{i}^{0} = X_{i}, X_{i}^{+1} = x_{i+1} \]

然后进行乘-加操作：

\[ Y = w_{1}X^{-1} + w_{2}X^{0} + w_{3}X^{1} \]

对于移位操作而言，其不涉及乘法运算，所以没有计算成本；TSM参考了这一特性，仅对视频帧进行移位操作，所以没有额外计算成本

缺陷

如果移位全部或者大多数通道，会出现以下两个问题：

  • 移位操作会导致数据移动，从而增加内存占用和推理延迟；
  • 移位操作会降低空间建模能力，从而导致模型性能退化。

设计

部分移位

论文测试了不同比例通道进行时间移位的效果。假定输入为8帧图像，执行以下测试：

  1. 没有移位操作；
  2. 部分移位操作：1/8、1/4、1/2；
  3. 全部移位操作

实验结果如上图（a）所示，执行部分移位操作能够有效减少额外产生的推理延迟时间

残差移位

为了保证空间建模能力，论文尝试了多种嵌入TSM的方式，如下图所示：

  • 最直接的方式是直接将TSM插入到每个卷积层或者残差块之前，如上图（a）所示，称之为就地移位（in-place shift）；
  • 论文同时测试了将TSM插入到残差块的残差分支中，如上图（b）所示，称之为残差移位（residual shift）。

因为恒等映射（identity mapping）的存在，网络仍旧能够有效的进行空间建模。实验结果如图2（b）所示：

  1. 残差移位能够有效的解决模型退化的问题；
  2. 使用1/4（每个方向1/8）的移位比例能够实现最佳的模型性能。

模块实现

  • 在训练过程中，可以使用双向移位（每个方向移位1/8通道）的TSM进行训练；
  • 在推理过程中，使用单向移位（仅向后移位1/8通道）。同时保留每帧图像在每个残差块计算中产生的前1/8特征图，作用于下一帧计算，如下图所示。

实验

训练设置

针对Kinetics数据集：

  1. 共100轮；
  2. 初始学习率为0.01（第40&80轮衰减1/10）；
  3. 权重衰减率为1e-4；
  4. 批量大小为64；
  5. 随机失活率为0.5

对于其他训练集，训练次数减半。针对预训练模型

  1. 对于大多数数据集，使用ImageNet预训练模型进行微调训练；
  2. 对于HMDB-51和UCF101，使用Kinetics预训练模型进行微调（冻结BN层）训练

实验结果

表一说明了相比于基准模型TSN，TSM在不同类型（是否依赖于时间关系）数据集均实现了性能提升

表三说明了在Kinetics数据集上，相比于基准模型TSN，在不同的Backbone中也能够实现性能提升

图五比较了TSM对于精确度和计算成本之间良好的平衡

表五说明了TSM拥有较低的计算延时和大吞吐性能

论文还比较了实时在线目标检测，相比于2D目标检测，其输入为单张图像，将嵌入TSM模块的ResNet作为Backbone，输入多帧数据，能够实现更好的检测性能

表八说明了TSM在边缘设备部署上拥有较低的延时