《TSM: Temporal Shift Module for Efficient Video Understanding》学习小记

《TSM: Temporal Shift Module for Efficient Video Understanding》学习小记

这里是这篇文章的链接，有兴趣的可以去下载原文进行阅读。
这里是文章相关的代码。
下面我主要是对原文的一些翻译和其中一些自己的观点，有其他意见的希望可以在评论区一起探讨。

作者的工作我认为有以下几点很突出的贡献：
1、实现在2D模型上达到3D模型的精度，极大的降低了计算！！！（因为3D模型需要极大的计算量，所以这对设备的性能是一个考验，对视频信息的实时性会有很大的挑战）
2、由于我们的backbone大部分是基于2D模型的，所以这样就方便了我们backbone的选择。
3、内存消耗低，缓存的仅仅是1/8的前一帧的特征，内存消耗极低。
4、实现了所有级别时间的融合。
本文提出的TSM网络主要解决的是传统的二维神经网络计算成本低，但不能捕获长期的时间关系;基于3D CNN的方法可以获得良好的性能，但计算量大，部署成本高的问题。所以这篇文章的作者提出了一个通用的、高效的时间移位模块(TSM)，该模块具有较高的效率和较高的性能。

特点
这个模块它可以实现3D CNN的性能，但保持了2D CNN相对较小的计算量。

技术手段
TSM在时间维度上移动了部分通道，从而促进相邻帧在信息之间的交流。而这一操作插入到2D CNN中又基本算是实现了在零计算量的情况下对时间信息的建模。文中还提到了他们的模型在侧重于时间建模的SomethingSomething-V1数据集上表现都是最好的。

图1 是TSM通过时间维度进行有效的时间建模，TSM有在线模式和离线模式的选择，在线模式的单向的时间移动，仅将过去的时间帧与当前帧进行融合（见图c），离线模式是双向的时间移动，将过去和将来的时间帧与当前帧融合在一起。其中不同的颜色块代表不同的时间T序列。
先来介绍一下3D模型的表示：R（N×C×T×H×W），其中N是批处理大小，C是通道数，T是时间维，H和W是空间分辨率（也就是图片的高和宽）。传统的2D CNN只能在尺寸T上单独运行，因此，没有时间建模会产生影响（图1a）。相反，我们的时间移位模块（TSM）沿时间维度向前和向后移动通道。如图1b所示，来自相邻帧的信息在移位后与当前帧混合在一起。作者的直觉是：卷积运算由移位和乘法累加组成。我们将时间维度偏移±1，并将乘积从时间维度折叠到通道维度。为了实时了解在线视频，无法将未来的信息转移到现在，因此我们使用单向TSM（图1c）进行在线视频理解。

这里有两个问题需要注意：这种shift操作会带来两个问题，（1）并不高效，尽管shift操作是基本zero FLOP的，但是这也同时会导致数据移动，这一步会使得延迟产生并且增加，而且尤其在视频文件中，该现象会更加显著，因为视频是5个维度的，上面说过；（2）并不准确，这时第二个问题，如果shifting太多的channels的话将会损坏原有帧的空间新，也就是移动了，那么就不完整了，并不包含所有的该张图片本身应该具有的信息！

作者的解决方案是：（1）改进的shifts策略：并不是shift所有的channels，而是只选择性的shift其中的一部分，该策略能够有效的减少数据移动所带来的时间复杂度；（2）TSM并不是直接被插入到从前往后的干道中的，而是以旁路的形式进行，如下图(b)，因此在获得了时序信息的同时不会对二维卷积的空间信息进行损害！（到这里基本上听起来就很靠谱啦！）

实现方法
我们首先考虑一个常规的卷积运算：为了简便起见，我们使用一维卷积，内核大小以3为例。假设卷积的权重为W =（w1，w2，w3），并且输入X是无限长的一维向量。卷积运算符Y = Conv（W，X）可以写为：Yi = w1X(i-1) + w2X(i) + w3X(i + 1)。 我们可以将卷积运算解耦为两步：移位和乘法累加：我们分别将输入X移位-1、0，+ 1并乘以w1，w2，w3，这些总和为Y。 正式地，移位操作为：


第一步移位可以免费执行，因为它只需要一个偏移地址指针。虽然第二步的计算开销更大，但是我们的时间移位模块将乘法累加合并到下面的2D卷积中，因此与基于2D CNN的模型相比，它不会带来额外的开销。

这是TSM框架的体系结构。视频被分割成N个大小相等的片段。每段采样一个帧。我们使用二维卷积从每一帧中提取空间特征。然后插入时间移位模块，实现无时间消耗的融合。
网络模型

1）Offline的双向TSM 模型
Backbone: ResNet-50
TSM的加入模式：Residual Block
特点1：对于每个插入的TSM模块，相当于其感受野得到了扩大2倍，因此可以获得很好的时序建模！
特点2：TSM可以应用的任意的现成的2DCNN上面，使其具有3DCNN的效果但是同时维持2DCNN的计算量
（2）Online的单向TSM模型
方式：把之前的frames shifts到当前的frames上，从而也达到了时序信息建模的作用，从而实现在线的recognition功能！
推理图如下：在推理的过程中，对于每一帧（比如Ft），只保存其首先的1/8的feature map。（作者经过消融实验发现1/8是一个最好的shift值）并且保存在cache中，对于下一帧(比如F t+1 ), 那么就是用Ft上的保存的1/8和当前帧Ft+1上的7/8组合起来。从而生成下一层feature map。如此循环下去！
利用这样的单向shift策略！
有以下特点：
特点1：推理过程中的低时延；
特点2：低内存消耗
特点3：多层时序信息的融合（网络的每一层都能得到融合！）

实验对比部分就不贴出来了，可以自行去查看文献。（总之就是很牛B！）
核心代码：
代码比较好懂，就不解释了。

class TemporalShift(nn.Module):
    def __init__(self, net, n_segment=3, n_div=8, inplace=False):
        super(TemporalShift, self).__init__()
        self.net = net
        self.n_segment = n_segment
        self.fold_div = n_div
        self.inplace = inplace
        if inplace:
            print('=> Using in-place shift...')
        print('=> Using fold div: {}'.format(self.fold_div))

    def forward(self, x):
        x = self.shift(x, self.n_segment, fold_div=self.fold_div, inplace=self.inplace)
        return self.net(x)

    @staticmethod
    def shift(x, n_segment, fold_div=3, inplace=False):
        nt, c, h, w = x.size()
        n_batch = nt // n_segment
        x = x.view(n_batch, n_segment, c, h, w)

        fold = c // fold_div
        if inplace:
            out = InplaceShift.apply(x, fold)
        else:
            out = torch.zeros_like(x)
            out[:, :-1, :fold] = x[:, 1:, :fold]  # shift left
            out[:, 1:, fold: 2 * fold] = x[:, :-1, fold: 2 * fold]  # shift right
            out[:, :, 2 * fold:] = x[:, :, 2 * fold:]  # not shift

        return out.view(nt, c, h, w)