TSM（Temporal Shift Module）源码解析

TSM（Temporal Shift Module）源码解析

论文名：TSM: Temporal Shift Module for Efficient Video Understanding

代码链接：https://github.com/mit-han-lab/temporal-shift-module

代码的主要结构如下：

python file	function
mian.py	主要的训练函数
opts.py	代码的参数配置
ops/dataset.py	数据集的载入部分，核心是__getitem__函数。
ops/dataset_config.py	用于配置不同的数据集
ops/models.py	组装模型
ops/temporal_shift.py	核心的temporal shift操作

1.opts.py是参数配置。除了路径、超参外，还有几个参数要注意（不同于TSN的点）：

parser.add_argument('--shift', default=False, action="store_true", help='use shift for models')
parser.add_argument('--shift_div', default=8, type=int, help='number of div for shift (default: 8)')

– modality 表示输入的类型，RGB/FLOW

–num_segments 表示每个视频采样frames数，一般是8或者16。

–shift 表示是否加入tsm模块。

–shift_div 表示shift的特征的比例，一般是8。表示2*1/8比例的特征会移动，其中1/8的特征做shift left, 另1/8的特征做shift right。

2.dataset_config.py是数据级的配置。

每个数据集实现一个return_xxx(modality)

返回数据集支持的子类名，train_list路径，val_list路径，数据集的根路径等信息。

3.dataset.py 是数据集的载入部分。

dataset.py的主要功能就是对数据集进行读取，并且对其稀疏采样，返回稀疏采样后得到的数据集。

dataset.py中实现了TSNDataSet类，处理原始数据，此类继承于torch.utils.data.dataset类。

其中首先定义了一个简单类：

3.1 VideoRecord，用于封装一个视频内容，包括图片的路径，frames的数量，标签信息。

class VideoRecord(object):
    def __init__(self, row):
        self._data = row

    @property
    def path(self):
        return self._data[0]

    @property
    def num_frames(self):
        return int(self._data[1])

    @property
    def label(self):
        return int(self._data[2])

3.2 TSNDataSet：

class TSNDataSet(data.Dataset):
    def __init__(self, root_path, list_file,
                 num_segments=3, new_length=1, modality='RGB',
                 image_tmpl='img_{:05d}.jpg', transform=None,
                 random_shift=True, test_mode=False,
                 remove_missing=False, dense_sample=False, twice_sample=False):

        self.root_path = root_path
        self.list_file = list_file
        self.num_segments = num_segments
        self.new_length = new_length
        self.modality = modality
        self.image_tmpl = image_tmpl
        self.transform = transform
        self.random_shift = random_shift
        self.test_mode = test_mode
        self.remove_missing = remove_missing
        self.dense_sample = dense_sample  # using dense sample as I3D
        self.twice_sample = twice_sample  # twice sample for more validation
        if self.dense_sample:
            print('=> Using dense sample for the dataset...')
        if self.twice_sample:
            print('=> Using twice sample for the dataset...')

        if self.modality == 'RGBDiff':
            self.new_length += 1  # Diff needs one more image to calculate diff

        self._parse_list()

设置一些参数和参数默认值之后，调用了_parse_list()函数，tmq是一个长度为训练数据数量的列表。每个值都是VIDEORecord对象，包含一个列表和3个属性，列表长度为3，用空格键分割，分别为帧路径、该视频含有多少帧和帧标签。然后调用VideoRecord()函数，将内容写入到一个VideoRecord的list中去。

 tmp = [x.strip().split(' ') for x in open(self.list_file)]

self.video_list = [VideoRecord(item) for item in tmp]

3.3 TSNDataSet需要实现核心函数__getitem()

首先我们需要获取视频对应的image和label，image的格式为[ n*t, c, h, w]

    def __getitem__(self, index):
        record = self.video_list[index]

那么对于一个视频，如何获得num_segments个帧呢？

        if not self.test_mode:  # test_mode: False; random_shift: True;
            segment_indices = self._sample_indices(record) if self.random_shift else self._get_val_indices(record)
        else:
            segment_indices = self._get_test_indices(record)
        return self.get(record, segment_indices)

有3种不同采样函数（针对于train、test、val）:

其中，_sample_indice函数是针对于train的采样方式；默认情况下， _sample_indice函数会随机采取num_segments个index，有稠密采样(dense)和稀疏采样(normal)两种方式.

_get_test_indices会稀疏且固定地得到num_segments个index。

简单举个例子，num_frames=120, num_segments=3的时候，

_sample_indices中normal sample会随机返回：[ 4，44，84]，[ 5，45，85]， [ 11，51，91]。

_sample_indices中dense sample会随机返回：[15, 36, 57], [30, 51, 72], [44, 65, 86]。

_get_test_indices中dense_sample同理

_get_test_indices中twice_sample会随机返回：[11,31,51,1,21,41]

最终，每个视频都采样num_segments个帧，getitem返回维度为：[n * t, c, h, w] （frames）， 1 （label）。

4.models.py

models.py的主要功能是对之后的训练模型做准备；首先使用一些经典的模型作为基础，如resnet50,针对不同的输入模态，对最后一层全连接层进行修改，得到我们TSN模型，而在其中又引入了是否加入TSM模块，从而得到我们所需要的TSM模型。

init函数设置一些参数和参数默认值,通过调用函数修改模型得到TSN模型，在调用TSM模块函数，从而得到TSM模型，其中init函数调用了

1 调用 _prepare_base_model(base_model)构建出基础的模型

2 调用_prepare_tsn(num_class)，用于根据不同数据集的子类数，适配fc层大小

3 对于flow和rgbdiff的输入，调用_construct_flow_model和_construct_diff_model更改第一个卷积核的大小

4.1 _prepare_base_model()函数

而调用make_temporal_shit()函数加入tsm模块：

    def _prepare_base_model(self, base_model):
        print('=> base model: {}'.format(base_model))

        if 'resnet' in base_model:
            # torchvision.models.resnet50
            self.base_model = getattr(torchvision.models, base_model)(True if self.pretrain == 'imagenet' else False)
            if self.is_shift:  # 默认false
                print('Adding temporal shift...')
                from ops.temporal_shift import make_temporal_shift
                make_temporal_shift(self.base_model, self.num_segments,
                                    n_div=self.shift_div, place=self.shift_place, temporal_pool=self.temporal_pool)

            if self.non_local:  # 默认False
                print('Adding non-local module...')
                from ops.non_local import make_non_local
                make_non_local(self.base_model, self.num_segments)

            self.base_model.last_layer_name = 'fc'
            self.input_size = 224
            self.input_mean = [0.485, 0.456, 0.406]
            self.input_std = [0.229, 0.224, 0.225]
		   # 1代表输出特征图的大小;torch.Size([2, 32, 16, 16])----->torch.Size([2, 32, 1, 1])
            self.base_model.avgpool = nn.AdaptiveAvgPool2d(1)  

4.2 _prepare_tsn函数

_prepare_tsn函数的功能在于对已知的basemodel网络结构进行修改，微调最后一层（全连接层）的结构，成为适合该数据集输出的形式。

    def _prepare_tsn(self, num_class):
        # 获取模型最后一层输入层的维度
        feature_dim = getattr(self.base_model, self.base_model.last_layer_name).in_features
        # 如果dropout==0，直接添加新的全连接层，输出维度是num_class
        if self.dropout == 0:
            setattr(self.base_model, self.base_model.last_layer_name, nn.Linear(feature_dim, num_class))
            self.new_fc = None
        # 如果有dropout！=0，添加dropout层，然后再添加全连接层，输出维度是num_class
        else:
            setattr(self.base_model, self.base_model.last_layer_name, nn.Dropout(p=self.dropout))
            self.new_fc = nn.Linear(feature_dim, num_class)

        std = 0.001
        if self.new_fc is None:
            normal_(getattr(self.base_model, self.base_model.last_layer_name).weight, 0, std)
            constant_(getattr(self.base_model, self.base_model.last_layer_name).bias, 0)
        else:
            if hasattr(self.new_fc, 'weight'):
                normal_(self.new_fc.weight, 0, std)
                constant_(self.new_fc.bias, 0)
        return feature_dim

4.3 make_temporal_shift()函数 TSM模块的核心

make_temporal_shift(self.base_model, self.num_segments,
                                    n_div=self.shift_div, place=self.shift_place, temporal_pool=self.temporal_pool)

那么TemporalShift(b, n_segment=this_segment, n_div=n_div又是怎么实现的呢？

对于[n*t, c, h, w]的输入，t是segment的值，首先reshape成[n, t, c, h, w]，

如果当前feature map的通道256，fold_div=8，那么有256/8的特征进行shift left，256/8的特征进行shift right。其他一部分的特征不动。

@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:
            # Due to some out of order error when performing parallel computing. 
            # May need to write a CUDA kernel.
            raise NotImplementedError  
            # 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)

举个简单的例子：

当c = 8, num_segment=4, 2维度的特征表示如下：

0_xx代表第一帧的特征。1_xx代码第二帧的特征，每个特征有8个通道。原始特征如下：

当fold_div = 8的时候，移动后如下：

可见第一帧中融入了第二帧的特征，第二帧中融入了第三帧和第二帧的特征。

当fold_div=4的时候，移动的部分会更多，即当前帧的特征中会包含更多前一帧和后一帧的信息。

5.main.py 是训练主函数

最后我们来讲解训练主函数，将上述将的类和函数串联起来。

5.1 调用dataset_config.return_dataset中获得各个路径的信息

5.2 实例化好模型、优化器

5.3 载入好预训练模型或从训练中一般恢复

5.4 准备好数据train_loader，val_loader

5.5 然后开始每个epoch迭代，保存一个最新的模型和最好的模型

adjust_learning_rate，根据策略调整每个epoch的学习率。

    for epoch in range(args.start_epoch, args.epochs):
        # Sets the learning rate to the initial LR decayed by 10 every 30 epochs
        adjust_learning_rate(optimizer, epoch, args.lr_type, args.lr_steps)

        # train for one epoch
        train(train_loader, model, criterion, optimizer, epoch, log_training, tf_writer)

        # evaluate on validation set
        if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
            prec1 = validate(val_loader, model, criterion, epoch, log_training, tf_writer)

            # remember best prec@1 and save checkpoint
            is_best = prec1 > best_prec1
            best_prec1 = max(prec1, best_prec1)
            tf_writer.add_scalar('acc/test_top1_best', best_prec1, epoch)

            output_best = 'Best Prec@1: %.3f\n' % (best_prec1)
            print(output_best)
            log_training.write(output_best + '\n')
            log_training.flush()  # 刷新缓冲区

            save_checkpoint({
                'epoch': epoch + 1,
                'arch': args.arch,
                'state_dict': model.state_dict(),
                'optimizer': optimizer.state_dict(),
                'best_prec1': best_prec1,
            }, is_best)

train的核心函数如下：

def train(train_loader, model, criterion, optimizer, epoch, log, tf_writer):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    if args.no_partialbn:  # 默认false
        model.module.partialBN(False)
    else:
        model.module.partialBN(True)

    # switch to train mode
    model.train()

    end = time.time()  # 返回当前时间戳
    for i, (input, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)
        raise RuntimeError

        input_var = torch.autograd.Variable(input)
        target_var = torch.autograd.Variable(target)

        # compute output
        output = model(input_var)
        loss = criterion(output, target_var)

        # print('********************')
        # print(loss)
        # raise RuntimeError

        # measure accuracy and record loss
        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
        losses.update(loss.item(), input.size(0))
        top1.update(prec1.item(), input.size(0))
        top5.update(prec5.item(), input.size(0))

        # compute gradient and do SGD step
        loss.backward()

        if args.clip_gradient is not None:  # None
            total_norm = clip_grad_norm_(model.parameters(), args.clip_gradient)

        optimizer.step()
        optimizer.zero_grad()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % args.print_freq == 0:   # print_freq = 20
            output = ('Epoch: [{0}][{1}/{2}], lr: {lr:.5f}\t'
                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                      'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                      'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                      'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                      'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                epoch, i, len(train_loader), batch_time=batch_time,
                data_time=data_time, loss=losses, top1=top1, top5=top5, lr=optimizer.param_groups[-1]['lr'] * 0.1))  # TODO
            print(output)
            log.write(output + '\n')
            log.flush()

    tf_writer.add_scalar('loss/train', losses.avg, epoch)
    tf_writer.add_scalar('acc/train_top1', top1.avg, epoch)
    tf_writer.add_scalar('acc/train_top5', top5.avg, epoch)
    tf_writer.add_scalar('lr', optimizer.param_groups[-1]['lr'], epoch)