duan8/tsm

Temporal Shift Module

TSM-R50 from "TSM: Temporal Shift Module for Efficient Video Understanding" https://arxiv.org/abs/1811.08383

TSM is a widely used Action Recognition model. This TensorRT implementation is tested with TensorRT 5.1 and TensorRT 7.2.

For the PyTorch implementation, you can refer to open-mmlab/mmaction2 or mit-han-lab/temporal-shift-module.

More details about the shift module(which is the core of TSM) could to test_shift.py.

Tutorial

• An example could refer to demo.sh

  • Requirements: Successfully installed torch>=1.3.0, torchvision

• Step 1: Train/Download TSM-R50 checkpoints from offical Github repo or MMAction2

  • Supported settings: num_segments, shift_div, num_classes.
  • Fixed settings: backbone(ResNet50), shift_place(blockres), temporal_pool(False).

• Step 2: Convert PyTorch checkpoints to TensorRT weights.

python gen_wts.py /path/to/pytorch.pth --out-filename /path/to/tensorrt.wts

• Step 3: Modify configs in tsm_r50.py.

BATCH_SIZE = 1
NUM_SEGMENTS = 8
INPUT_H = 224
INPUT_W = 224
OUTPUT_SIZE = 400
SHIFT_DIV = 8

• Step 4: Inference with tsm_r50.py.

usage: tsm_r50.py [-h] [--tensorrt-weights TENSORRT_WEIGHTS] [--input-video INPUT_VIDEO] [--save-engine-path SAVE_ENGINE_PATH] [--load-engine-path LOAD_ENGINE_PATH] [--test-mmaction2] [--mmaction2-config MMACTION2_CONFIG] [--mmaction2-checkpoint MMACTION2_CHECKPOINT]

optional arguments:
  -h, --help            show this help message and exit
  --tensorrt-weights TENSORRT_WEIGHTS
                        Path to TensorRT weights, which is generated by gen_weights.py
  --input-video INPUT_VIDEO
                        Path to local video file
  --save-engine-path SAVE_ENGINE_PATH
                        Save engine to local file
  --load-engine-path LOAD_ENGINE_PATH
                        Saved engine file path
  --test-mmaction2      Compare TensorRT results with MMAction2 Results
  --mmaction2-config MMACTION2_CONFIG
                        Path to MMAction2 config file
  --mmaction2-checkpoint MMACTION2_CHECKPOINT
                        Path to MMAction2 checkpoint url or file path

TODO

• Python Shift module.
• Generate wts of official tsm and mmaction2 tsm.
• Python API Definition
• Test with mmaction2 demo
• Tutorial
• C++ API Definition