mmaction2 中 TSM（Temporal Shift Module）视频动作识别实战：原理、模型库与训练测试指南

• 人工智能
• 计算机视觉
• 深度学习

TSM（Temporal Shift Module）是 mmaction2 中"以 2D CNN 的成本获得 3D CNN 级时序建模能力"的代表性算法。本文围绕仓库中 configs/recognition/tsm/README.md 展开，完整覆盖 TSM 的模型库（Kinetics-400 与 Something-Something V2 的全部预设模型）、采样策略（1x1x8 / 1x1x16 / dense / Non-Local 变体）、训练与测试命令，并结合 ResNetTSM 源码、TSMHead 与 TSM 专用优化器构造器 深入讲解 Temporal Shift 的底层实现。读完后，你可以独立复现 TSM 的训练/测试流程，理解每个配置项（num_segments、shift_div、shift_place、fc_lr5等）如何映射到底层行为，并能完成 MobileOne 骨干的 reparameterization 部署。

1. TSM 的核心思想：零参数、零算力的时序建模

TSM 论文（Lin et al., ICCV 2019）针对的核心矛盾是：2D CNN 计算便宜但无法建模帧间关系，3D CNN 效果好但部署代价高。TSM 的解法是在特征图的通道维上做"移位"——把部分通道沿时间维度平移，使相邻帧之间交换信息。论文摘要中的关键结论是：

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.

也就是说，TSM不新增任何参数、不增加任何卷积计算，仅通过张量的 reshape 与拼接就实现了时序建模。这正是模型库表中 TSM 各变体 FLOPs 与参数量与"纯 2D 骨干"完全一致（如 ResNet50 均为 32.88G / 23.87M，仅取决于采样帧数而非结构）的原因。

1.1 Temporal Shift 的源码实现

核心实现在 mmaction/models/backbones/resnet_tsm.py 的TemporalShift类中。其shift静态方法的工作流程是：

1. 输入特征形状为[N, C, H, W]，其中N = batch × num_segments（每个视频采样出的num_segments帧被展平到 batch 维）；
2. 将Nreshape 为[N // num_segments, num_segments, C, H*W]，使"帧"成为独立的一个维度；
3. 按fold = C // shift_div把通道切成三段left_split / mid_split / right_split：
   • left_split（前 1/8 通道）：沿时间维向左移 1 帧（取后段 + 补 1 个空帧）；
   • mid_split（中间 1/8 通道）：沿时间维向右移 1 帧（补 1 个空帧 + 取前段）；
   • right_split（其余 6/8 通道）：不做任何移位；
4. 三段拼回[N, C, H, W]后送入被包裹的原始网络层self.net。

几个值得注意的实现细节：

• 移位用left_split - left_split构造空帧而非torch.zeros_like，源码注释说明这是因为 Caffe 推理后端的数组必须由计算产生（resnet_tsm.py#L103-L110）；
• shift_div=8是论文默认值，控制参与移位的通道比例（左右各 1/8）；
• forward是"先 shift 再算原层"，因此 TSM 本质是一个包装器（wrapper），可套在任意 2D 层外面。

1.2 shift 的插入位置：shift_place

ResNetTSM通过shift_place参数决定 shift 插在哪里（resnet_tsm.py#L188-L253）：

• 'block'：把整个残差块包进TemporalShift；
• 'blockres'（默认值）：只包住每个残差块的conv1.conv（1×1 卷积），且仅当 stage 内块数满足i % n_round == 0时插入。其中n_round默认 1，当layer3子模块数 ≥ 23（即 ResNet-101 的 36 个 block 场景）时自动变为 2，即每隔一个块才插入一次 shift——这是 TSM 论文对深骨干减半 shift 频率的对应实现。

此外ResNetTSM还内置了两个可选扩展：

• Non-Local：non_local=(0,0,0,0)按 stage 指定在哪些残差块后插入 3D Non-Local 模块，由NL3DWrapper将 2D 特征转回 5D 时序张量做 Non-Local 再展平（resnet_tsm.py#L13-L41）。模型库中三个nl-*配置就来自这里；
• Temporal Pooling：temporal_pool=True时在 layer2 外包一层MaxPool3d(kernel=(3,1,1), stride=(2,1,1))，使 segment 数减半，从而在更少的帧上运行后部网络。

2. 模型结构：Recognizer2D + ResNetTSM + TSMHead

2.1 基础配置

所有 TSM 配置都继承基础模型配置 configs/base/models/tsm_r50.py：

preprocess_cfg = dict( mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]) model = dict( type='Recognizer2D', backbone=dict( type='ResNetTSM', pretrained='torchvision://resnet50', depth=50, norm_eval=False, shift_div=8), cls_head=dict( type='TSMHead', num_classes=400, in_channels=2048, spatial_type='avg', consensus=dict(type='AvgConsensus', dim=1), dropout_ratio=0.5, init_std=0.001, is_shift=True, average_clips='prob'), data_preprocessor=dict(type='ActionDataPreprocessor', **preprocess_cfg), train_cfg=None, test_cfg=None)

要点：

• Recognizer2D是 2D 识别器，pretrained='torchvision://resnet50'表示用 ImageNet 预训练的 2D 权重初始化骨干——ResNetTSM.load_original_weights会专门处理.net/.block这类 wrapper 前缀，把 torchvision 的 key 正确映射到被 TSM 包裹后的模块名（resnet_tsm.py#L311-L360）；
• shift_div=8即通道三分中每段占 1/8；
• TSMHead的consensus=dict(type='AvgConsensus', dim=1)表示对num_segments个时间片段取平均共识，average_clips='prob'表示对每个 clip 先算 softmax 概率再平均；
• in_channels=2048对应 ResNet50 最后一层特征，换骨干时需同步修改（如 MobileNetV2 为 1280）。

2.2 TSMHead 的前向逻辑

TSMHead 的前向过程：

1. 输入[N*num_segs, in_channels, H, W]，做自适应全局平均池化 → flatten → dropout → 全连接，得到每帧独立分类得分；
2. 将得分view回[N, num_segments, num_classes]（若启用 temporal_pooling 则为num_segments // 2）；
3. 经AvgConsensus在 segment 维取平均，得到[N, num_classes]最终得分。

值得注意num_segments是head 与 backbone 共享的超参数：采样 16 帧的 1x1x16 配置中，backbone 与 cls_head 必须同时写num_segments=16（见 tsm_imagenet-pretrained-r50_8xb16-1x1x16-50e_kinetics400-rgb.py 第 4 行model = dict(backbone=dict(num_segments=16), cls_head=dict(num_segments=16))）。

2.3 TSM 专属的学习率策略

所有 TSM 训练配置都使用TSMOptimWrapperConstructor（tsm_optim_wrapper_constructor.py），它对参数分组施加论文中的学习率约定：

这就是配置中paramwise_cfg=dict(fc_lr5=True)的作用。以 Kinetics-400 主配置 为例，完整优化器设置为 SGD（lr=0.02, momentum=0.9, weight_decay=0.0001）+ 梯度裁剪（max_norm=20），学习率调度为前 5 epoch 线性 warmup（start_factor=0.1）+ 25/45 epoch 两个 milestone 各乘 0.1。

3. 帧采样策略：如何读懂 "1x1x8" 之类的采样记号

配置名中的采样记号c x len x num（clip 数 × clip 长度 × 片段数）直接映射到数据管线的SampleFrames：

# 1x1x8：8 个 clip，每个 clip 取 1 帧 dict(type='SampleFrames', clip_len=1, frame_interval=1, num_clips=8) # 1x1x16：16 个 clip，每个 clip 取 1 帧 dict(type='SampleFrames', clip_len=1, frame_interval=1, num_clips=16) # dense 变体：均匀连续采样 dict(type='DenseSampleFrames', clip_len=1, frame_interval=1, num_clips=8)

三条关键规则：

1. num_clips必须与num_segments一致：1x1x8 对应num_segments=8，1x1x16 对应num_segments=16，否则TemporalShift无法正确地把 batch 维还原成[N, num_segments, C, H, W]；
2. 训练增强统一为MultiScaleCrop（input_size=224, scales=(1, 0.875, 0.75, 0.66), num_fixed_crops=13）+ 50% 水平翻转；
3. 测试裁剪决定表中的 "x 10 crop / x 3 crop"：Kinetics 用TenCrop（10 裁剪），Something-Something V2 用ThreeCrop（3 裁剪），STHV2 的SampleFrames还启用了twice_sample=True。

STHV2 配置（tsm_imagenet-pretrained-r50_8xb16-1x1x8-50e_sthv2-rgb.py）还有一个数据集特有细节：翻转增强携带flip_label_map={86: 87, 87: 86, 93: 94, 94: 93, 166: 167, 167: 166}，因为 STHV2 存在水平镜像后含义互换的动作对（如 "pushing from left to right" 与 "pushing from right to left"），num_classes也因此是 174 而非 400。

4. 模型库：预设模型与结果

4.1 Kinetics-400

4.2 Something-Something V2

上述模型的官方 checkpoint 与训练日志均随模型库一并发布（以 mmaction v1.0 模型库为准），配置文件中auto_scale_lr.base_batch_size=128即"8 GPU × 16 样本"的原始 batch size。

4.3 使用注意事项（原文档说明）

1. GPU 数变化时的学习率：表中 GPU 数为产出 checkpoint 所用卡数。若使用不同的 GPU 数或每卡视频数，建议在调用tools/train.py时加--auto-scale-lr，它会按实际 batch size 与原始 batch size 的比例自动缩放学习率；
2. Kinetics400 验证集：官方使用的 Kinetics400 验证集由 19796 个视频组成，对应的数据列表（每行格式为video_id, num_frames, label_index）与类别映射表可在官方模型库的 dataset 下载区获取；
3. MobileOne 骨干支持推理重参数化：可用 tools/convert/reparameterize_model.py 转换 checkpoint，并切换到部署配置 tsm_imagenet-pretrained-mobileone-s4_deploy_8xb16-1x1x16-50e_kinetics400-rgb.py。

关于数据准备的完整说明，可参考 Kinetics400 数据准备文档。

5. 训练

使用标准训练入口：

python tools/train.py ${CONFIG_FILE} [optional arguments]

原文档给出的确定性训练示例（在 Kinetics-400 上训练 TSM-R50，固定随机种子）：

python tools/train.py configs/recognition/tsm/tsm_imagenet-pretrained-r50_8xb16-1x1x8-50e_kinetics400-rgb.py \ --seed=0 --deterministic

以该 50e 配置为例，一次训练的完整训练侧参数如下，可直接作为自定义数据集的改造模板：

• 数据：VideoDataset，训练列表kinetics400_train_list_videos.txt，视频根目录data/kinetics400/videos_train，file_client_args=dict(io_backend='disk')；
• batch：batch_size=16（每卡）× 8 卡 = 128，num_workers=8, persistent_workers=True；
• 训练循环：train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=50, val_begin=1, val_interval=1)；
• 调度器：LinearLR（0→5 epoch，start_factor=0.1）+MultiStepLR（milestones=[25, 45], gamma=0.1）；
• 优化器：TSMOptimWrapperConstructor+fc_lr5=True+ SGD(lr=0.02, momentum=0.9, weight_decay=1e-4)，梯度裁剪max_norm=20；
• 检查点：default_hooks = dict(checkpoint=dict(interval=3, max_keep_ckpts=3))；
• 评估：val_evaluator = dict(type='AccMetric')。

更多训练参数细节可参考 Training and Test Tutorial 中的Training部分。

6. 测试

python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} [optional arguments]

原文档示例（在 Kinetics-400 上测试并把逐视频结果 dump 到 pkl 文件）：

python tools/test.py configs/recognition/tsm/tsm_imagenet-pretrained-r50_8xb16-1x1x8-50e_kinetics400-rgb.py \ checkpoints/SOME_CHECKPOINT.pth --dump result.pkl

测试流程走test_dataloader：batch_size=1（10-crop 推理下每个视频展开为 80 帧图像，小 batch 避免显存溢出）、test_pipeline使用SampleFrames(..., test_mode=True)+TenCrop/ThreeCrop，最终由AccMetric汇总 Top-1/Top-5。更多参数可参考 Training and Test Tutorial 的Test部分。

7. MobileOne 骨干的 Reparameterization 部署

MobileOne 骨干训练时是"多分支"结构，推理时可把分支合并为单卷积以提升效率。mmaction2 中的实现链路是：

1. MobileOneTSM 继承 mmpretrain 的MobileOne，其make_temporal_shift与 ResNet 版不同——把TemporalShift套在整个 block外层（shift_place等价于'block'），源码注释明确说明这是为了让 reparameterization 可以正常工作（若像 ResNet 版那样包在conv1内层，则无法安全合并分支）；
2. 训练配置 tsm_imagenet-pretrained-mobileone-s4_8xb16-1x1x16-50e_kinetics400-rgb.py 中该模型auto_scale_lr默认开启（enable=True），适合不同卡数直接复用；
3. 用 reparameterize_model.py 执行model.backbone.switch_to_deploy()并保存新权重，部署配置只需在训练配置上追加一行model = dict(backbone=dict(deploy=True))（见 deploy 配置）。

命令形式为：

python tools/convert/reparameterize_model.py ${CONFIG} ${CHECKPOINT} ${SAVE_PATH}

8. 引用

如果使用了 TSM 模型，请引用：

@inproceedings{lin2019tsm, title={TSM: Temporal Shift Module for Efficient Video Understanding}, author={Lin, Ji and Gan, Chuang and Han, Song}, booktitle={Proceedings of the IEEE International Conference on Computer Vision}, year={2019} }

使用 Non-Local 变体（nl-dot-product/nl-gaussian/nl-embedded-gaussian三个配置，通过 backbone 的non_local与non_local_cfg参数启用）时，还应引用：

@article{Nonlocal2018, author = {Xiaolong Wang and Ross Girshick and Abhinav Gupta and Kaiming He}, title = {Non-local Neural Networks}, journal = {CVPR}, year = {2018} }

• 人工智能
• 计算机视觉
• 深度学习